Merge "Cancel back invoke when window focus has lost" into tm-qpr-dev am: f52b49886c

Original change: https://googleplex-android-review.googlesource.com/c/platform/frameworks/base/+/19415501

Change-Id: I7d4fc2e6f3df820a4a4594b0368ff6bfbb12820f
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>

12 files changed, 132 insertions, 129 deletions

diff --git a/core/java/android/app/IActivityTaskManager.aidl b/core/java/android/app/IActivityTaskManager.aidl
index dc6825c85705..0bf16a0c471b 100644
--- a/core/java/android/app/IActivityTaskManager.aidl
+++ b/core/java/android/app/IActivityTaskManager.aidl
@@ -69,6 +69,7 @@ import android.os.WorkSource;
 import android.service.voice.IVoiceInteractionSession;
 import android.view.IRecentsAnimationRunner;
 import android.view.IRemoteAnimationRunner;
+import android.view.IWindowFocusObserver;
 import android.view.RemoteAnimationDefinition;
 import android.view.RemoteAnimationAdapter;
 import android.window.IWindowOrganizerController;
@@ -352,6 +353,8 @@ interface IActivityTaskManager {
      * Prepare the back navigation in the server. This setups the leashed for sysui to animate
      * the back gesture and returns the data needed for the animation.
      * @param requestAnimation true if the caller wishes to animate the back navigation
+     * @param focusObserver a remote callback to nofify shell when the focused window lost focus.
      */
-    android.window.BackNavigationInfo startBackNavigation(in boolean requestAnimation);
+    android.window.BackNavigationInfo startBackNavigation(in boolean requestAnimation,
+            in IWindowFocusObserver focusObserver);
 }
diff --git a/core/java/android/view/ViewRootImpl.java b/core/java/android/view/ViewRootImpl.java
index 3347163a1683..9319b41a5b5d 100644
--- a/core/java/android/view/ViewRootImpl.java
+++ b/core/java/android/view/ViewRootImpl.java
@@ -3761,7 +3761,6 @@ public final class ViewRootImpl implements ViewParent,
             }
         }
         mFirstInputStage.onWindowFocusChanged(hasWindowFocus);
-        mOnBackInvokedDispatcher.onWindowFocusChanged(hasWindowFocus);
 
         // NOTE: there's no view visibility (appeared / disapparead) events when the windows focus
         // is lost, so we don't need to to force a flush - there might be other events such as
diff --git a/core/java/android/window/BackNavigationInfo.java b/core/java/android/window/BackNavigationInfo.java
index 0ab6db58ab4b..dd4901417671 100644
--- a/core/java/android/window/BackNavigationInfo.java
+++ b/core/java/android/window/BackNavigationInfo.java
@@ -118,7 +118,7 @@ public final class BackNavigationInfo implements Parcelable {
      *                                back preview.
      * @param onBackInvokedCallback   The back callback registered by the current top level window.
      */
-    public BackNavigationInfo(@BackTargetType int type,
+    private BackNavigationInfo(@BackTargetType int type,
             @Nullable RemoteAnimationTarget departingAnimationTarget,
             @Nullable SurfaceControl screenshotSurface,
             @Nullable HardwareBuffer screenshotBuffer,
diff --git a/core/java/android/window/ImeOnBackInvokedDispatcher.java b/core/java/android/window/ImeOnBackInvokedDispatcher.java
index f1a052b61c59..14f52282b3aa 100644
--- a/core/java/android/window/ImeOnBackInvokedDispatcher.java
+++ b/core/java/android/window/ImeOnBackInvokedDispatcher.java
@@ -83,8 +83,7 @@ public class ImeOnBackInvokedDispatcher implements OnBackInvokedDispatcher, Parc
         final Bundle bundle = new Bundle();
         // Always invoke back for ime without checking the window focus.
         final IOnBackInvokedCallback iCallback =
-                new WindowOnBackInvokedDispatcher.OnBackInvokedCallbackWrapper(callback,
-                        () -> true);
+                new WindowOnBackInvokedDispatcher.OnBackInvokedCallbackWrapper(callback);
         bundle.putBinder(RESULT_KEY_CALLBACK, iCallback.asBinder());
         bundle.putInt(RESULT_KEY_PRIORITY, priority);
         bundle.putInt(RESULT_KEY_ID, callback.hashCode());
diff --git a/core/java/android/window/WindowOnBackInvokedDispatcher.java b/core/java/android/window/WindowOnBackInvokedDispatcher.java
index 02c5ebcc184e..1d396be6b478 100644
--- a/core/java/android/window/WindowOnBackInvokedDispatcher.java
+++ b/core/java/android/window/WindowOnBackInvokedDispatcher.java
@@ -32,7 +32,6 @@ import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Objects;
 import java.util.TreeMap;
-import java.util.function.Supplier;
 
 /**
  * Provides window based implementation of {@link OnBackInvokedDispatcher}.
@@ -65,7 +64,6 @@ public class WindowOnBackInvokedDispatcher implements OnBackInvokedDispatcher {
     private final TreeMap<Integer, ArrayList<OnBackInvokedCallback>>
             mOnBackInvokedCallbacks = new TreeMap<>();
     private final Checker mChecker;
-    private boolean mHasFocus;
 
     public WindowOnBackInvokedDispatcher(boolean applicationCallBackEnabled) {
         mChecker = new Checker(applicationCallBackEnabled);
@@ -191,7 +189,7 @@ public class WindowOnBackInvokedDispatcher implements OnBackInvokedDispatcher {
                                     .ImeOnBackInvokedCallback
                                 ? ((ImeOnBackInvokedDispatcher.ImeOnBackInvokedCallback)
                                         callback).getIOnBackInvokedCallback()
-                                : new OnBackInvokedCallbackWrapper(callback, this::hasFocus);
+                                : new OnBackInvokedCallbackWrapper(callback);
                 callbackInfo = new OnBackInvokedCallbackInfo(iCallback, priority);
             }
             mWindowSession.setOnBackInvokedCallbackInfo(mWindow, callbackInfo);
@@ -200,17 +198,6 @@ public class WindowOnBackInvokedDispatcher implements OnBackInvokedDispatcher {
         }
     }
 
-    /**
-     * Called when window focus changed.
-     */
-    public void onWindowFocusChanged(boolean hasFocus) {
-        mHasFocus = hasFocus;
-    }
-
-    private boolean hasFocus() {
-        return mHasFocus;
-    }
-
     public OnBackInvokedCallback getTopCallback() {
         if (mAllCallbacks.isEmpty()) {
             return null;
@@ -234,11 +221,8 @@ public class WindowOnBackInvokedDispatcher implements OnBackInvokedDispatcher {
 
     static class OnBackInvokedCallbackWrapper extends IOnBackInvokedCallback.Stub {
         private final WeakReference<OnBackInvokedCallback> mCallback;
-        private final Supplier<Boolean> mHasFocus;
-        OnBackInvokedCallbackWrapper(@NonNull OnBackInvokedCallback callback,
-                @NonNull Supplier<Boolean> hasFocus) {
+        OnBackInvokedCallbackWrapper(@NonNull OnBackInvokedCallback callback) {
             mCallback = new WeakReference<>(callback);
-            mHasFocus = hasFocus;
         }
 
         @Override
@@ -278,10 +262,6 @@ public class WindowOnBackInvokedDispatcher implements OnBackInvokedDispatcher {
                 if (callback == null) {
                     return;
                 }
-                if (!mHasFocus.get()) {
-                    Log.w(TAG, "Skip back invoke due to current focus has lost.");
-                    return;
-                }
                 callback.onBackInvoked();
             });
         }
diff --git a/core/tests/coretests/src/android/window/BackNavigationTest.java b/core/tests/coretests/src/android/window/BackNavigationTest.java
index ce69f12a065b..bbbc4230903a 100644
--- a/core/tests/coretests/src/android/window/BackNavigationTest.java
+++ b/core/tests/coretests/src/android/window/BackNavigationTest.java
@@ -91,7 +91,8 @@ public class BackNavigationTest {
     private void assertCallbackIsCalled(CountDownLatch latch) {
         try {
             mInstrumentation.getUiAutomation().waitForIdle(500, 1000);
-            BackNavigationInfo info = ActivityTaskManager.getService().startBackNavigation(true);
+            BackNavigationInfo info = ActivityTaskManager.getService()
+                    .startBackNavigation(true, null);
             assertNotNull("BackNavigationInfo is null", info);
             assertNotNull("OnBackInvokedCallback is null", info.getOnBackInvokedCallback());
             info.getOnBackInvokedCallback().onBackInvoked();
diff --git a/core/tests/coretests/src/android/window/WindowOnBackInvokedDispatcherTest.java b/core/tests/coretests/src/android/window/WindowOnBackInvokedDispatcherTest.java
index b5194f637395..f448cb3091e7 100644
--- a/core/tests/coretests/src/android/window/WindowOnBackInvokedDispatcherTest.java
+++ b/core/tests/coretests/src/android/window/WindowOnBackInvokedDispatcherTest.java
@@ -66,7 +66,6 @@ public class WindowOnBackInvokedDispatcherTest {
         MockitoAnnotations.initMocks(this);
         mDispatcher = new WindowOnBackInvokedDispatcher(true /* applicationCallbackEnabled */);
         mDispatcher.attachToWindow(mWindowSession, mWindow);
-        mDispatcher.onWindowFocusChanged(true);
     }
 
     private void waitForIdle() {
@@ -153,31 +152,4 @@ public class WindowOnBackInvokedDispatcherTest {
         waitForIdle();
         verify(mCallback2).onBackStarted();
     }
-
-    @Test
-    public void skipBackInvokeWhenNoFocus() throws RemoteException {
-        ArgumentCaptor<OnBackInvokedCallbackInfo> captor =
-                ArgumentCaptor.forClass(OnBackInvokedCallbackInfo.class);
-
-        mDispatcher.registerOnBackInvokedCallback(
-                OnBackInvokedDispatcher.PRIORITY_DEFAULT, mCallback1);
-
-        verify(mWindowSession, times(1)).setOnBackInvokedCallbackInfo(
-                Mockito.eq(mWindow),
-                captor.capture());
-
-        verify(mWindowSession).setOnBackInvokedCallbackInfo(Mockito.eq(mWindow), captor.capture());
-
-        // Should invoke back if it's still in focused.
-        captor.getValue().getCallback().onBackInvoked();
-        waitForIdle();
-        verify(mCallback1).onBackInvoked();
-
-        // In case the focus has lost during back gesture.
-        mDispatcher.onWindowFocusChanged(false);
-
-        captor.getValue().getCallback().onBackInvoked();
-        waitForIdle();
-        verifyZeroInteractions(mCallback1);
-    }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationController.java
index d53a98cf5b07..02211220822a 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationController.java
@@ -32,11 +32,13 @@ import android.graphics.PointF;
 import android.hardware.HardwareBuffer;
 import android.net.Uri;
 import android.os.Handler;
+import android.os.IBinder;
 import android.os.RemoteException;
 import android.os.SystemProperties;
 import android.os.UserHandle;
 import android.provider.Settings.Global;
 import android.util.Log;
+import android.view.IWindowFocusObserver;
 import android.view.MotionEvent;
 import android.view.RemoteAnimationTarget;
 import android.view.SurfaceControl;
@@ -107,6 +109,24 @@ public class BackAnimationController implements RemoteCallable<BackAnimationCont
         mTransitionInProgress = false;
     };
 
+    @VisibleForTesting
+    final IWindowFocusObserver mFocusObserver = new IWindowFocusObserver.Stub() {
+        @Override
+        public void focusGained(IBinder inputToken) { }
+        @Override
+        public void focusLost(IBinder inputToken) {
+            mShellExecutor.execute(() -> {
+                if (!mBackGestureStarted || mTransitionInProgress) {
+                    // If an uninterruptible transition is already in progress, we should ignore
+                    // this due to the transition may cause focus lost. (alpha = 0)
+                    return;
+                }
+                setTriggerBack(false);
+                onGestureFinished(false);
+            });
+        }
+    };
+
     public BackAnimationController(
             @NonNull @ShellMainThread ShellExecutor shellExecutor,
             @NonNull @ShellBackgroundThread Handler backgroundHandler,
@@ -266,17 +286,17 @@ public class BackAnimationController implements RemoteCallable<BackAnimationCont
                 // Let the animation initialized here to make sure the onPointerDownOutsideFocus
                 // could be happened when ACTION_DOWN, it may change the current focus that we
                 // would access it when startBackNavigation.
-                initAnimation(touchX, touchY);
+                onGestureStarted(touchX, touchY);
             }
             onMove(touchX, touchY, swipeEdge);
         } else if (keyAction == MotionEvent.ACTION_UP || keyAction == MotionEvent.ACTION_CANCEL) {
             ProtoLog.d(WM_SHELL_BACK_PREVIEW,
                     "Finishing gesture with event action: %d", keyAction);
-            onGestureFinished();
+            onGestureFinished(true);
         }
     }
 
-    private void initAnimation(float touchX, float touchY) {
+    private void onGestureStarted(float touchX, float touchY) {
         ProtoLog.d(WM_SHELL_BACK_PREVIEW, "initAnimation mMotionStarted=%b", mBackGestureStarted);
         if (mBackGestureStarted || mBackNavigationInfo != null) {
             Log.e(TAG, "Animation is being initialized but is already started.");
@@ -288,7 +308,8 @@ public class BackAnimationController implements RemoteCallable<BackAnimationCont
 
         try {
             boolean requestAnimation = mEnableAnimations.get();
-            mBackNavigationInfo = mActivityTaskManager.startBackNavigation(requestAnimation);
+            mBackNavigationInfo =
+                    mActivityTaskManager.startBackNavigation(requestAnimation, mFocusObserver);
             onBackNavigationInfoReceived(mBackNavigationInfo);
         } catch (RemoteException remoteException) {
             Log.e(TAG, "Failed to initAnimation", remoteException);
@@ -376,11 +397,18 @@ public class BackAnimationController implements RemoteCallable<BackAnimationCont
         dispatchOnBackProgressed(targetCallback, backEvent);
     }
 
-    private void onGestureFinished() {
+    private void onGestureFinished(boolean fromTouch) {
         ProtoLog.d(WM_SHELL_BACK_PREVIEW, "onGestureFinished() mTriggerBack == %s", mTriggerBack);
-        if (!mBackGestureStarted || mBackNavigationInfo == null) {
+        if (fromTouch) {
+            // Let touch reset the flag otherwise it will start a new back navigation and refresh
+            // the info when received a new move event.
+            mBackGestureStarted = false;
+        }
+
+        if (mTransitionInProgress || mBackNavigationInfo == null) {
             return;
         }
+
         int backType = mBackNavigationInfo.getType();
         boolean shouldDispatchToLauncher = shouldDispatchToLauncher(backType);
         IOnBackInvokedCallback targetCallback = shouldDispatchToLauncher
@@ -470,7 +498,6 @@ public class BackAnimationController implements RemoteCallable<BackAnimationCont
 
     private void finishAnimation() {
         ProtoLog.d(WM_SHELL_BACK_PREVIEW, "BackAnimationController: finishAnimation()");
-        mBackGestureStarted = false;
         mTouchEventDelta.set(0, 0);
         mInitTouchLocation.set(0, 0);
         BackNavigationInfo backNavigationInfo = mBackNavigationInfo;
@@ -480,6 +507,7 @@ public class BackAnimationController implements RemoteCallable<BackAnimationCont
         if (backNavigationInfo == null) {
             return;
         }
+
         RemoteAnimationTarget animationTarget = backNavigationInfo.getDepartingAnimationTarget();
         if (animationTarget != null) {
             if (animationTarget.leash != null && animationTarget.leash.isValid()) {
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/back/BackAnimationControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/back/BackAnimationControllerTest.java
index 31e55e7116f0..ba81602054a8 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/back/BackAnimationControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/back/BackAnimationControllerTest.java
@@ -20,6 +20,7 @@ import static android.window.BackNavigationInfo.KEY_TRIGGER_BACK;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.eq;
@@ -39,6 +40,7 @@ import android.graphics.Point;
 import android.graphics.Rect;
 import android.hardware.HardwareBuffer;
 import android.os.Handler;
+import android.os.IBinder;
 import android.os.RemoteCallback;
 import android.os.RemoteException;
 import android.provider.Settings;
@@ -121,24 +123,22 @@ public class BackAnimationControllerTest extends ShellTestCase {
             HardwareBuffer hardwareBuffer,
             int backType,
             IOnBackInvokedCallback onBackInvokedCallback) {
-        BackNavigationInfo navigationInfo = new BackNavigationInfo(
-                backType,
-                topAnimationTarget,
-                screenshotSurface,
-                hardwareBuffer,
-                new WindowConfiguration(),
-                new RemoteCallback((bundle) -> {}),
-                onBackInvokedCallback);
-        try {
-            doReturn(navigationInfo).when(mActivityTaskManager).startBackNavigation(anyBoolean());
-        } catch (RemoteException ex) {
-            ex.rethrowFromSystemServer();
-        }
+        BackNavigationInfo.Builder builder = new BackNavigationInfo.Builder()
+                .setType(backType)
+                .setDepartingAnimationTarget(topAnimationTarget)
+                .setScreenshotSurface(screenshotSurface)
+                .setScreenshotBuffer(hardwareBuffer)
+                .setTaskWindowConfiguration(new WindowConfiguration())
+                .setOnBackNavigationDone(new RemoteCallback((bundle) -> {}))
+                .setOnBackInvokedCallback(onBackInvokedCallback);
+
+        createNavigationInfo(builder);
     }
 
     private void createNavigationInfo(BackNavigationInfo.Builder builder) {
         try {
-            doReturn(builder.build()).when(mActivityTaskManager).startBackNavigation(anyBoolean());
+            doReturn(builder.build()).when(mActivityTaskManager)
+                    .startBackNavigation(anyBoolean(), any());
         } catch (RemoteException ex) {
             ex.rethrowFromSystemServer();
         }
@@ -297,6 +297,34 @@ public class BackAnimationControllerTest extends ShellTestCase {
         verify(mIOnBackInvokedCallback).onBackStarted();
     }
 
+
+    @Test
+    public void cancelBackInvokeWhenLostFocus() throws RemoteException {
+        mController.setBackToLauncherCallback(mIOnBackInvokedCallback);
+        RemoteAnimationTarget animationTarget = createAnimationTarget();
+
+        createNavigationInfo(animationTarget, null, null,
+                BackNavigationInfo.TYPE_RETURN_TO_HOME, null);
+
+        doMotionEvent(MotionEvent.ACTION_DOWN, 0);
+        // Check that back start and progress is dispatched when first move.
+        doMotionEvent(MotionEvent.ACTION_MOVE, 100);
+        verify(mIOnBackInvokedCallback).onBackStarted();
+
+        // Check that back invocation is dispatched.
+        mController.setTriggerBack(true);   // Fake trigger back
+
+        // In case the focus has been changed.
+        IBinder token = mock(IBinder.class);
+        mController.mFocusObserver.focusLost(token);
+        mShellExecutor.flushAll();
+        verify(mIOnBackInvokedCallback).onBackCancelled();
+
+        // No more back invoke.
+        doMotionEvent(MotionEvent.ACTION_UP, 0);
+        verify(mIOnBackInvokedCallback, never()).onBackInvoked();
+    }
+
     private void doMotionEvent(int actionDown, int coordinate) {
         mController.onMotionEvent(
                 coordinate, coordinate,
diff --git a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
index 9d8e0877584f..0260b78f49e4 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
@@ -222,6 +222,7 @@ import android.util.SparseArray;
 import android.util.TimeUtils;
 import android.util.proto.ProtoOutputStream;
 import android.view.IRecentsAnimationRunner;
+import android.view.IWindowFocusObserver;
 import android.view.RemoteAnimationAdapter;
 import android.view.RemoteAnimationDefinition;
 import android.view.WindowManager;
@@ -1027,7 +1028,7 @@ public class ActivityTaskManagerService extends IActivityTaskManager.Stub {
             mTaskSupervisor.setWindowManager(wm);
             mRootWindowContainer.setWindowManager(wm);
             if (mBackNavigationController != null) {
-                mBackNavigationController.setTaskSnapshotController(wm.mTaskSnapshotController);
+                mBackNavigationController.setWindowManager(wm);
             }
         }
     }
@@ -1833,13 +1834,14 @@ public class ActivityTaskManagerService extends IActivityTaskManager.Stub {
     }
 
     @Override
-    public BackNavigationInfo startBackNavigation(boolean requestAnimation) {
+    public BackNavigationInfo startBackNavigation(boolean requestAnimation,
+            IWindowFocusObserver observer) {
         mAmInternal.enforceCallingPermission(START_TASKS_FROM_RECENTS,
                 "startBackNavigation()");
         if (mBackNavigationController == null) {
             return null;
         }
-        return mBackNavigationController.startBackNavigation(mWindowManager, requestAnimation);
+        return mBackNavigationController.startBackNavigation(requestAnimation, observer);
     }
 
     /**
diff --git a/services/core/java/com/android/server/wm/BackNavigationController.java b/services/core/java/com/android/server/wm/BackNavigationController.java
index 9295c18fe80e..f967cf995e67 100644
--- a/services/core/java/com/android/server/wm/BackNavigationController.java
+++ b/services/core/java/com/android/server/wm/BackNavigationController.java
@@ -31,6 +31,7 @@ import android.os.RemoteCallback;
 import android.os.RemoteException;
 import android.os.SystemProperties;
 import android.util.Slog;
+import android.view.IWindowFocusObserver;
 import android.view.RemoteAnimationTarget;
 import android.view.SurfaceControl;
 import android.window.BackNavigationInfo;
@@ -45,10 +46,9 @@ import com.android.server.LocalServices;
  * Controller to handle actions related to the back gesture on the server side.
  */
 class BackNavigationController {
-
     private static final String TAG = "BackNavigationController";
-    @Nullable
-    private TaskSnapshotController mTaskSnapshotController;
+    private WindowManagerService mWindowManagerService;
+    private IWindowFocusObserver mFocusObserver;
 
     /**
      * Returns true if the back predictability feature is enabled
@@ -69,26 +69,13 @@ class BackNavigationController {
      * for the animation, or null if we don't know how to animate the current window and need to
      * fallback on dispatching the key event.
      */
-    @Nullable
-    BackNavigationInfo startBackNavigation(@NonNull WindowManagerService wmService,
-            boolean requestAnimation) {
-        return startBackNavigation(wmService, null, requestAnimation);
-    }
-
-    /**
-     * @param tx, a transaction to be used for the attaching the animation leash.
-     *            This is used in tests. If null, the object will be initialized with a new {@link
-     *            SurfaceControl.Transaction}
-     * @see #startBackNavigation(WindowManagerService, boolean)
-     */
     @VisibleForTesting
     @Nullable
-    BackNavigationInfo startBackNavigation(WindowManagerService wmService,
-            @Nullable SurfaceControl.Transaction tx, boolean requestAnimation) {
-
-        if (tx == null) {
-            tx = new SurfaceControl.Transaction();
-        }
+    BackNavigationInfo startBackNavigation(boolean requestAnimation,
+            IWindowFocusObserver observer) {
+        final WindowManagerService wmService = mWindowManagerService;
+        final SurfaceControl.Transaction tx = wmService.mTransactionFactory.get();
+        mFocusObserver = observer;
 
         int backType = BackNavigationInfo.TYPE_UNDEFINED;
 
@@ -110,13 +97,14 @@ class BackNavigationController {
         SurfaceControl animationLeashParent = null;
         HardwareBuffer screenshotBuffer = null;
         RemoteAnimationTarget topAppTarget = null;
+        WindowState window;
+
         int prevTaskId;
         int prevUserId;
         boolean prepareAnimation;
 
         BackNavigationInfo.Builder infoBuilder = new BackNavigationInfo.Builder();
         synchronized (wmService.mGlobalLock) {
-            WindowState window;
             WindowConfiguration taskWindowConfiguration;
             WindowManagerInternal windowManagerInternal =
                     LocalServices.getService(WindowManagerInternal.class);
@@ -189,6 +177,9 @@ class BackNavigationController {
                     backType = BackNavigationInfo.TYPE_CALLBACK;
                 }
                 infoBuilder.setOnBackInvokedCallback(callbackInfo.getCallback());
+                if (mFocusObserver != null) {
+                    window.registerFocusObserver(mFocusObserver);
+                }
             }
 
             ProtoLog.d(WM_DEBUG_BACK_PREVIEW, "startBackNavigation currentTask=%s, "
@@ -209,9 +200,13 @@ class BackNavigationController {
                     || currentActivity == null
                     || currentTask == null
                     || currentActivity.isActivityTypeHome()) {
-                return infoBuilder
-                        .setType(backType)
-                        .build();
+                infoBuilder.setType(BackNavigationInfo.TYPE_CALLBACK);
+                final WindowState finalFocusedWindow = window;
+                infoBuilder.setOnBackNavigationDone(new RemoteCallback(result ->
+                        onBackNavigationDone(result, finalFocusedWindow, finalFocusedWindow,
+                                BackNavigationInfo.TYPE_CALLBACK, null, null, false)));
+
+                return infoBuilder.setType(backType).build();
             }
 
             // We don't have an application callback, let's find the destination of the back gesture
@@ -345,11 +340,12 @@ class BackNavigationController {
             }
 
             int finalBackType = backType;
-            ActivityRecord finalprevActivity = prevActivity;
-            Task finalTask = currentTask;
+            final ActivityRecord finalprevActivity = prevActivity;
+            final Task finalTask = currentTask;
+            final WindowState finalFocusedWindow = window;
             RemoteCallback onBackNavigationDone = new RemoteCallback(result -> onBackNavigationDone(
-                    result, finalRemovedWindowContainer, finalBackType, finalTask,
-                    finalprevActivity, prepareAnimation));
+                    result, finalFocusedWindow, finalRemovedWindowContainer, finalBackType,
+                    finalTask, finalprevActivity, prepareAnimation));
             infoBuilder.setOnBackNavigationDone(onBackNavigationDone);
         }
 
@@ -382,8 +378,8 @@ class BackNavigationController {
     }
 
     private void onBackNavigationDone(
-            Bundle result, WindowContainer<?> windowContainer, int backType,
-            Task task, ActivityRecord prevActivity, boolean prepareAnimation) {
+            Bundle result, WindowState focusedWindow, WindowContainer<?> windowContainer,
+            int backType, Task task, ActivityRecord prevActivity, boolean prepareAnimation) {
         SurfaceControl surfaceControl = windowContainer.getSurfaceControl();
         boolean triggerBack = result != null && result.getBoolean(
                 BackNavigationInfo.KEY_TRIGGER_BACK);
@@ -412,6 +408,11 @@ class BackNavigationController {
             task.mBackGestureStarted = false;
         }
         resetSurfaces(windowContainer);
+
+        if (mFocusObserver != null) {
+            focusedWindow.unregisterFocusObserver(mFocusObserver);
+            mFocusObserver = null;
+        }
     }
 
     private HardwareBuffer getActivitySnapshot(@NonNull Task task,
@@ -425,10 +426,10 @@ class BackNavigationController {
     }
 
     private HardwareBuffer getTaskSnapshot(int taskId, int userId) {
-        if (mTaskSnapshotController == null) {
+        if (mWindowManagerService.mTaskSnapshotController == null) {
             return null;
         }
-        TaskSnapshot snapshot = mTaskSnapshotController.getSnapshot(taskId,
+        TaskSnapshot snapshot = mWindowManagerService.mTaskSnapshotController.getSnapshot(taskId,
                 userId, true /* restoreFromDisk */, false  /* isLowResolution */);
         return snapshot != null ? snapshot.getHardwareBuffer() : null;
     }
@@ -458,7 +459,7 @@ class BackNavigationController {
         }
     }
 
-    void setTaskSnapshotController(@Nullable TaskSnapshotController taskSnapshotController) {
-        mTaskSnapshotController = taskSnapshotController;
+    void setWindowManager(WindowManagerService wm) {
+        mWindowManagerService = wm;
     }
 }
diff --git a/services/tests/wmtests/src/com/android/server/wm/BackNavigationControllerTests.java b/services/tests/wmtests/src/com/android/server/wm/BackNavigationControllerTests.java
index fc41a94d3355..c2ca0a227f26 100644
--- a/services/tests/wmtests/src/com/android/server/wm/BackNavigationControllerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/BackNavigationControllerTests.java
@@ -29,12 +29,9 @@ import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.doAnswer;
 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 static org.mockito.Mockito.when;
 
 import android.annotation.NonNull;
@@ -42,7 +39,6 @@ import android.annotation.Nullable;
 import android.hardware.HardwareBuffer;
 import android.os.RemoteException;
 import android.platform.test.annotations.Presubmit;
-import android.view.SurfaceControl;
 import android.view.WindowManager;
 import android.window.BackEvent;
 import android.window.BackNavigationInfo;
@@ -75,8 +71,7 @@ public class BackNavigationControllerTests extends WindowTestsBase {
         LocalServices.removeServiceForTest(WindowManagerInternal.class);
         mWindowManagerInternal = mock(WindowManagerInternal.class);
         LocalServices.addService(WindowManagerInternal.class, mWindowManagerInternal);
-        TaskSnapshotController taskSnapshotController = createMockTaskSnapshotController();
-        mBackNavigationController.setTaskSnapshotController(taskSnapshotController);
+        mBackNavigationController.setWindowManager(mWm);
     }
 
     @Test
@@ -84,19 +79,14 @@ public class BackNavigationControllerTests extends WindowTestsBase {
         Task task = createTopTaskWithActivity();
         IOnBackInvokedCallback callback = withSystemCallback(task);
 
-        SurfaceControl.Transaction tx = mock(SurfaceControl.Transaction.class);
-        BackNavigationInfo backNavigationInfo = mBackNavigationController.startBackNavigation(mWm,
-                tx, true);
+        BackNavigationInfo backNavigationInfo =
+                mBackNavigationController.startBackNavigation(true, null);
         assertWithMessage("BackNavigationInfo").that(backNavigationInfo).isNotNull();
         assertThat(backNavigationInfo.getDepartingAnimationTarget()).isNotNull();
         assertThat(backNavigationInfo.getTaskWindowConfiguration()).isNotNull();
         assertThat(backNavigationInfo.getOnBackInvokedCallback()).isEqualTo(callback);
         assertThat(typeToString(backNavigationInfo.getType()))
                 .isEqualTo(typeToString(BackNavigationInfo.TYPE_RETURN_TO_HOME));
-
-        verify(tx, atLeastOnce()).apply();
-        verify(tx, times(1)).reparent(any(),
-                eq(backNavigationInfo.getDepartingAnimationTarget().leash));
     }
 
     @Test
@@ -243,7 +233,7 @@ public class BackNavigationControllerTests extends WindowTestsBase {
 
     @Nullable
     private BackNavigationInfo startBackNavigation() {
-        return mBackNavigationController.startBackNavigation(mWm, new StubTransaction(), true);
+        return mBackNavigationController.startBackNavigation(true, null);
     }
 
     @NonNull