Create OnBackInvokedCallbackInfo to wrap callback and its priority.

Main motivation is to store a back callback's exact priority value in WM. This is required by the IME migration (ag/17076160) to compare the priority levels of IME window callback and focused window callback in BackNavigationController.

This also consolidates the WindowState#mSystemOnBackInvokedCallback and WindowState#mApplicationOnBackInvokedCallback fields into one field, as tracking two fields for one callback was error prone. We had to remember to clear the application / system field when the other field is set, and failing to do so has resulted in bugs such as b/222675481.


Bug: 224856664
Test: atest BackNavigationControllerTest
Test: atest WindowOnBackInvokedDispatcherTest
Test: m -j and test back behavior throughout the system on apps that
opted in and out.

Change-Id: Ic57113610d934f33d2c9ca4cef59f39a9b87e832

core/java/android/view/IWindowSession.aidl

@@ -37,7 +37,7 @@ import android.view.Surface;
 import android.view.SurfaceControl;
 import android.view.SurfaceControl.Transaction;
 import android.window.ClientWindowFrames;
-import android.window.IOnBackInvokedCallback;
+import android.window.OnBackInvokedCallbackInfo;
 
 import java.util.List;
 
@@ -371,14 +371,14 @@ interface IWindowSession {
             in String hashAlgorithm, in RemoteCallback callback);
 
     /**
-     * Sets the {@link IOnBackInvokedCallback} to be invoked for a window when back is triggered.
+     * Sets the {@link OnBackInvokedCallbackInfo} containing the callback to be invoked for
+     * a window when back is triggered.
      *
      * @param window The token for the window to set the callback to.
-     * @param callback The {@link IOnBackInvokedCallback} to set.
+     * @param callbackInfo The {@link OnBackInvokedCallbackInfo} to set.
-     * @param priority The priority of the callback.
      */
-    oneway void setOnBackInvokedCallback(
+    oneway void setOnBackInvokedCallbackInfo(
-            IWindow window, IOnBackInvokedCallback callback, int priority);
+            IWindow window, in OnBackInvokedCallbackInfo callbackInfo);
 
     /**
      * Clears a touchable region set by {@link #setInsets}.

core/java/android/view/WindowlessWindowManager.java

@@ -31,7 +31,7 @@ import android.os.RemoteException;
 import android.util.Log;
 import android.util.MergedConfiguration;
 import android.window.ClientWindowFrames;
-import android.window.IOnBackInvokedCallback;
+import android.window.OnBackInvokedCallbackInfo;
 
 import java.util.HashMap;
 import java.util.List;
@@ -529,8 +529,8 @@ public class WindowlessWindowManager implements IWindowSession {
     }
 
     @Override
-    public void setOnBackInvokedCallback(IWindow iWindow,
+    public void setOnBackInvokedCallbackInfo(IWindow iWindow,
-            IOnBackInvokedCallback iOnBackInvokedCallback, int priority) throws RemoteException { }
+            OnBackInvokedCallbackInfo callbackInfo) throws RemoteException { }
 
     @Override
     public boolean dropForAccessibility(IWindow window, int x, int y) {

core/java/android/window/OnBackInvokedCallbackInfo.aidl

@@ -0,0 +1,21 @@
+/*
+ * Copyright (C) 2021 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 android.window;
+
+/** @hide */
+parcelable OnBackInvokedCallbackInfo;

core/java/android/window/OnBackInvokedCallbackInfo.java

@@ -0,0 +1,85 @@
+/*
+ * 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 android.window;
+
+import android.annotation.NonNull;
+import android.os.Parcel;
+import android.os.Parcelable;
+
+/**
+ * Data object to hold an {@link IOnBackInvokedCallback} and its priority.
+ * @hide
+ */
+public final class OnBackInvokedCallbackInfo implements Parcelable {
+    @NonNull
+    private final IOnBackInvokedCallback mCallback;
+    private @OnBackInvokedDispatcher.Priority int mPriority;
+
+    public OnBackInvokedCallbackInfo(@NonNull IOnBackInvokedCallback callback, int priority) {
+        mCallback = callback;
+        mPriority = priority;
+    }
+
+    private OnBackInvokedCallbackInfo(@NonNull Parcel in) {
+        mCallback = IOnBackInvokedCallback.Stub.asInterface(in.readStrongBinder());
+        mPriority = in.readInt();
+    }
+
+    @Override
+    public int describeContents() {
+        return 0;
+    }
+
+    @Override
+    public void writeToParcel(@NonNull Parcel dest, int flags) {
+        dest.writeStrongInterface(mCallback);
+        dest.writeInt(mPriority);
+    }
+
+    public static final Creator<OnBackInvokedCallbackInfo> CREATOR =
+            new Creator<OnBackInvokedCallbackInfo>() {
+                @Override
+                public OnBackInvokedCallbackInfo createFromParcel(Parcel in) {
+                    return new OnBackInvokedCallbackInfo(in);
+                }
+
+                @Override
+                public OnBackInvokedCallbackInfo[] newArray(int size) {
+                    return new OnBackInvokedCallbackInfo[size];
+                }
+            };
+
+    public boolean isSystemCallback() {
+        return mPriority == OnBackInvokedDispatcher.PRIORITY_SYSTEM;
+    }
+
+    @NonNull
+    public IOnBackInvokedCallback getCallback() {
+        return mCallback;
+    }
+
+    @OnBackInvokedDispatcher.Priority
+    public int getPriority() {
+        return mPriority;
+    }
+
+    @Override
+    public String toString() {
+        return "OnBackInvokedCallbackInfo{"
+                + "mCallback=" + mCallback + ", mPriority=" + mPriority + '}';
+    }
+}

core/java/android/window/WindowOnBackInvokedDispatcher.java

@@ -160,11 +160,12 @@ public class WindowOnBackInvokedDispatcher implements OnBackInvokedDispatcher {
         }
         try {
             if (callback == null) {
-                mWindowSession.setOnBackInvokedCallback(mWindow, null, PRIORITY_DEFAULT);
+                mWindowSession.setOnBackInvokedCallbackInfo(mWindow, null);
             } else {
                 int priority = mAllCallbacks.get(callback);
-                mWindowSession.setOnBackInvokedCallback(
+                mWindowSession.setOnBackInvokedCallbackInfo(
-                        mWindow, new OnBackInvokedCallbackWrapper(callback), priority);
+                        mWindow, new OnBackInvokedCallbackInfo(
+                                new OnBackInvokedCallbackWrapper(callback), priority));
             }
             if (DEBUG && callback == null) {
                 Log.d(TAG, TextUtils.formatSimple("setTopOnBackInvokedCallback(null) Callers:%s",

core/tests/coretests/src/android/window/WindowOnBackInvokedDispatcherTest.java

@@ -16,6 +16,7 @@
 
 package android.window;
 
+import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.isNull;
 import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.times;
@@ -73,24 +74,23 @@ public class WindowOnBackInvokedDispatcherTest {
 
     @Test
     public void propagatesTopCallback_samePriority() throws RemoteException {
-        ArgumentCaptor<IOnBackInvokedCallback> captor =
+        ArgumentCaptor<OnBackInvokedCallbackInfo> captor =
-                ArgumentCaptor.forClass(IOnBackInvokedCallback.class);
+                ArgumentCaptor.forClass(OnBackInvokedCallbackInfo.class);
 
         mDispatcher.registerOnBackInvokedCallback(
                 OnBackInvokedDispatcher.PRIORITY_DEFAULT, mCallback1);
         mDispatcher.registerOnBackInvokedCallback(
                 OnBackInvokedDispatcher.PRIORITY_DEFAULT, mCallback2);
 
-        verify(mWindowSession, times(2)).setOnBackInvokedCallback(
+        verify(mWindowSession, times(2)).setOnBackInvokedCallbackInfo(
                 Mockito.eq(mWindow),
-                captor.capture(),
+                captor.capture());
-                Mockito.eq(OnBackInvokedDispatcher.PRIORITY_DEFAULT));
+        captor.getAllValues().get(0).getCallback().onBackStarted();
-        captor.getAllValues().get(0).onBackStarted();
         waitForIdle();
         verify(mCallback1).onBackStarted();
         verifyZeroInteractions(mCallback2);
 
-        captor.getAllValues().get(1).onBackStarted();
+        captor.getAllValues().get(1).getCallback().onBackStarted();
         waitForIdle();
         verify(mCallback2).onBackStarted();
         verifyNoMoreInteractions(mCallback1);
@@ -98,19 +98,19 @@ public class WindowOnBackInvokedDispatcherTest {
 
     @Test
     public void propagatesTopCallback_differentPriority() throws RemoteException {
-        ArgumentCaptor<IOnBackInvokedCallback> captor =
+        ArgumentCaptor<OnBackInvokedCallbackInfo> captor =
-                ArgumentCaptor.forClass(IOnBackInvokedCallback.class);
+                ArgumentCaptor.forClass(OnBackInvokedCallbackInfo.class);
 
         mDispatcher.registerOnBackInvokedCallback(
                 OnBackInvokedDispatcher.PRIORITY_OVERLAY, mCallback1);
         mDispatcher.registerOnBackInvokedCallback(
                 OnBackInvokedDispatcher.PRIORITY_DEFAULT, mCallback2);
 
-        verify(mWindowSession).setOnBackInvokedCallback(
+        verify(mWindowSession).setOnBackInvokedCallbackInfo(
-                Mockito.eq(mWindow), captor.capture(),
+                Mockito.eq(mWindow), captor.capture());
-                Mockito.eq(OnBackInvokedDispatcher.PRIORITY_OVERLAY));
         verifyNoMoreInteractions(mWindowSession);
-        captor.getValue().onBackStarted();
+        assertEquals(captor.getValue().getPriority(), OnBackInvokedDispatcher.PRIORITY_OVERLAY);
+        captor.getValue().getCallback().onBackStarted();
         waitForIdle();
         verify(mCallback1).onBackStarted();
     }
@@ -127,17 +127,14 @@ public class WindowOnBackInvokedDispatcherTest {
         verifyZeroInteractions(mWindowSession);
 
         mDispatcher.unregisterOnBackInvokedCallback(mCallback2);
-        verify(mWindowSession).setOnBackInvokedCallback(
+        verify(mWindowSession).setOnBackInvokedCallbackInfo(Mockito.eq(mWindow), isNull());
-                Mockito.eq(mWindow),
-                isNull(),
-                Mockito.eq(OnBackInvokedDispatcher.PRIORITY_DEFAULT));
     }
 
 
     @Test
     public void propagatesTopCallback_sameInstanceAddedTwice() throws RemoteException {
-        ArgumentCaptor<IOnBackInvokedCallback> captor =
+        ArgumentCaptor<OnBackInvokedCallbackInfo> captor =
-                ArgumentCaptor.forClass(IOnBackInvokedCallback.class);
+                ArgumentCaptor.forClass(OnBackInvokedCallbackInfo.class);
 
         mDispatcher.registerOnBackInvokedCallback(OnBackInvokedDispatcher.PRIORITY_OVERLAY,
                 mCallback1
@@ -150,11 +147,8 @@ public class WindowOnBackInvokedDispatcherTest {
         reset(mWindowSession);
         mDispatcher.registerOnBackInvokedCallback(
                 OnBackInvokedDispatcher.PRIORITY_OVERLAY, mCallback2);
-        verify(mWindowSession).setOnBackInvokedCallback(
+        verify(mWindowSession).setOnBackInvokedCallbackInfo(Mockito.eq(mWindow), captor.capture());
-                Mockito.eq(mWindow),
+        captor.getValue().getCallback().onBackStarted();
-                captor.capture(),
-                Mockito.eq(OnBackInvokedDispatcher.PRIORITY_OVERLAY));
-        captor.getValue().onBackStarted();
         waitForIdle();
         verify(mCallback2).onBackStarted();
     }

data/etc/services.core.protolog.json

@@ -1351,12 +1351,6 @@
       "group": "WM_DEBUG_FOCUS_LIGHT",
       "at": "com\/android\/server\/wm\/ActivityRecord.java"
     },
-    "-767349300": {
-      "message": "%s: Setting back callback %s. Client IWindow %s",
-      "level": "DEBUG",
-      "group": "WM_DEBUG_BACK_PREVIEW",
-      "at": "com\/android\/server\/wm\/WindowState.java"
-    },
     "-766059044": {
       "message": "Display id=%d selected orientation %s (%d), got rotation %s (%d)",
       "level": "VERBOSE",
@@ -1909,6 +1903,12 @@
       "group": "WM_DEBUG_SYNC_ENGINE",
       "at": "com\/android\/server\/wm\/BLASTSyncEngine.java"
     },
+    "-228813488": {
+      "message": "%s: Setting back callback %s",
+      "level": "DEBUG",
+      "group": "WM_DEBUG_BACK_PREVIEW",
+      "at": "com\/android\/server\/wm\/WindowState.java"
+    },
     "-208825711": {
       "message": "shouldWaitAnimatingExit: isWallpaperTarget: %s",
       "level": "DEBUG",
@@ -2377,12 +2377,6 @@
       "group": "WM_DEBUG_REMOTE_ANIMATIONS",
       "at": "com\/android\/server\/wm\/RemoteAnimationController.java"
     },
-    "250620778": {
-      "message": "startBackNavigation task=%s, topRunningActivity=%s, applicationBackCallback=%s, systemBackCallback=%s, currentFocus=%s",
-      "level": "DEBUG",
-      "group": "WM_DEBUG_BACK_PREVIEW",
-      "at": "com\/android\/server\/wm\/BackNavigationController.java"
-    },
     "251812577": {
       "message": "Register display organizer=%s uid=%d",
       "level": "VERBOSE",
@@ -2785,12 +2779,6 @@
       "group": "WM_DEBUG_APP_TRANSITIONS",
       "at": "com\/android\/server\/wm\/WindowState.java"
     },
-    "599897753": {
-      "message": "Previous Activity is %s. Back type is %s",
-      "level": "DEBUG",
-      "group": "WM_DEBUG_BACK_PREVIEW",
-      "at": "com\/android\/server\/wm\/BackNavigationController.java"
-    },
     "600140673": {
       "message": "checkBootAnimationComplete: Waiting for anim complete",
       "level": "INFO",
@@ -3061,6 +3049,12 @@
       "group": "WM_DEBUG_REMOTE_ANIMATIONS",
       "at": "com\/android\/server\/wm\/RemoteAnimationController.java"
     },
+    "878005951": {
+      "message": "startBackNavigation task=%s, topRunningActivity=%s, callbackInfo=%s, currentFocus=%s",
+      "level": "DEBUG",
+      "group": "WM_DEBUG_BACK_PREVIEW",
+      "at": "com\/android\/server\/wm\/BackNavigationController.java"
+    },
     "892244061": {
       "message": "Waiting for drawn %s: removed=%b visible=%b mHasSurface=%b drawState=%d",
       "level": "INFO",

services/core/java/com/android/server/wm/BackNavigationController.java

@@ -34,7 +34,7 @@ import android.util.Slog;
 import android.view.RemoteAnimationTarget;
 import android.view.SurfaceControl;
 import android.window.BackNavigationInfo;
-import android.window.IOnBackInvokedCallback;
+import android.window.OnBackInvokedCallbackInfo;
 import android.window.TaskSnapshot;
 
 import com.android.internal.annotations.VisibleForTesting;
@@ -158,36 +158,30 @@ class BackNavigationController {
             }
 
             // Now let's find if this window has a callback from the client side.
-            IOnBackInvokedCallback applicationCallback = null;
+            OnBackInvokedCallbackInfo callbackInfo = null;
-            IOnBackInvokedCallback systemCallback = null;
             if (window != null) {
                 activityRecord = window.mActivityRecord;
                 task = window.getTask();
-                applicationCallback = window.getApplicationOnBackInvokedCallback();
+                callbackInfo = window.getOnBackInvokedCallbackInfo();
-                if (applicationCallback != null) {
+                if (callbackInfo == null) {
-                    backType = BackNavigationInfo.TYPE_CALLBACK;
+                    Slog.e(TAG, "No callback registered, returning null.");
-                    infoBuilder.setOnBackInvokedCallback(applicationCallback);
+                    return null;
-                } else {
-                    systemCallback = window.getSystemOnBackInvokedCallback();
-                    infoBuilder.setOnBackInvokedCallback(systemCallback);
                 }
+                if (!callbackInfo.isSystemCallback()) {
+                    backType = BackNavigationInfo.TYPE_CALLBACK;
+                }
+                infoBuilder.setOnBackInvokedCallback(callbackInfo.getCallback());
             }
 
             ProtoLog.d(WM_DEBUG_BACK_PREVIEW, "startBackNavigation task=%s, "
-                            + "topRunningActivity=%s, applicationBackCallback=%s, "
+                            + "topRunningActivity=%s, callbackInfo=%s, currentFocus=%s",
-                            + "systemBackCallback=%s, currentFocus=%s",
+                    task, activityRecord, callbackInfo, window);
-                    task, activityRecord, applicationCallback, systemCallback, window);
 
             if (window == null) {
                 Slog.e(TAG, "Window is null, returning null.");
                 return null;
             }
 
-            if (systemCallback == null && applicationCallback == null) {
-                Slog.e(TAG, "No callback registered, returning null.");
-                return null;
-            }
-
             // If we don't need to set up the animation, we return early. This is the case when
             // - We have an application callback.
             // - We don't have any ActivityRecord or Task to animate.

services/core/java/com/android/server/wm/Session.java

@@ -74,8 +74,7 @@ import android.view.SurfaceSession;
 import android.view.View;
 import android.view.WindowManager;
 import android.window.ClientWindowFrames;
-import android.window.IOnBackInvokedCallback;
+import android.window.OnBackInvokedCallbackInfo;
-import android.window.OnBackInvokedDispatcher;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.os.logging.MetricsLoggerWrapper;
@@ -937,18 +936,16 @@ class Session extends IWindowSession.Stub implements IBinder.DeathRecipient {
     }
 
     @Override
-    public void setOnBackInvokedCallback(
+    public void setOnBackInvokedCallbackInfo(
             IWindow window,
-            IOnBackInvokedCallback onBackInvokedCallback,
+            OnBackInvokedCallbackInfo callbackInfo) {
-            @OnBackInvokedDispatcher.Priority int priority) {
         synchronized (mService.mGlobalLock) {
             WindowState windowState = mService.windowForClientLocked(this, window, false);
             if (windowState == null) {
                 Slog.e(TAG_WM,
-                        "setOnBackInvokedCallback(): No window state for package:"
+                        "setOnBackInvokedCallback(): No window state for package:" + mPackageName);
-                                + mPackageName);
             } else {
-                windowState.setOnBackInvokedCallback(onBackInvokedCallback, priority);
+                windowState.setOnBackInvokedCallbackInfo(callbackInfo);
             }
         }
     }

services/core/java/com/android/server/wm/WindowState.java

@@ -249,8 +249,7 @@ import android.view.animation.Animation;
 import android.view.animation.AnimationUtils;
 import android.view.animation.Interpolator;
 import android.window.ClientWindowFrames;
-import android.window.IOnBackInvokedCallback;
+import android.window.OnBackInvokedCallbackInfo;
-import android.window.WindowOnBackInvokedDispatcher;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.policy.KeyInterceptionInfo;
@@ -821,12 +820,9 @@ class WindowState extends WindowContainer<WindowState> implements WindowManagerP
     };
 
     /**
-     * @see #setOnBackInvokedCallback(IOnBackInvokedCallback)
+     * @see #setOnBackInvokedCallbackInfo(OnBackInvokedCallbackInfo)
      */
-    // TODO(b/224856664): Consolidate application and system callback into one.
+    private OnBackInvokedCallbackInfo mOnBackInvokedCallbackInfo;
-    private IOnBackInvokedCallback mApplicationOnBackInvokedCallback;
-    private IOnBackInvokedCallback mSystemOnBackInvokedCallback;
-
     @Override
     WindowState asWindowState() {
         return this;
@@ -1083,28 +1079,16 @@ class WindowState extends WindowContainer<WindowState> implements WindowManagerP
      * called when a back navigation action is initiated.
      * @see BackNavigationController
      */
-    void setOnBackInvokedCallback(
+    void setOnBackInvokedCallbackInfo(
-            @Nullable IOnBackInvokedCallback onBackInvokedCallback, int priority) {
+            @Nullable OnBackInvokedCallbackInfo callbackInfo) {
-        ProtoLog.d(WM_DEBUG_BACK_PREVIEW, "WindowState: Setting back callback %s (priority: %d) "
+        ProtoLog.d(WM_DEBUG_BACK_PREVIEW, "%s: Setting back callback %s",
-                        + "(Client IWindow: %s). (WindowState: %s)",
+                this, callbackInfo);
-                onBackInvokedCallback, priority, mClient, this);
+        mOnBackInvokedCallbackInfo = callbackInfo;
-        if (priority >= WindowOnBackInvokedDispatcher.PRIORITY_DEFAULT) {
-            mApplicationOnBackInvokedCallback = onBackInvokedCallback;
-            mSystemOnBackInvokedCallback = null;
-        } else {
-            mApplicationOnBackInvokedCallback = null;
-            mSystemOnBackInvokedCallback = onBackInvokedCallback;
-        }
     }
 
     @Nullable
-    IOnBackInvokedCallback getApplicationOnBackInvokedCallback() {
+    OnBackInvokedCallbackInfo getOnBackInvokedCallbackInfo() {
-        return mApplicationOnBackInvokedCallback;
+        return mOnBackInvokedCallbackInfo;
-    }
-
-    @Nullable
-    IOnBackInvokedCallback getSystemOnBackInvokedCallback() {
-        return mSystemOnBackInvokedCallback;
     }
 
     interface PowerManagerWrapper {
@@ -2496,8 +2480,7 @@ class WindowState extends WindowContainer<WindowState> implements WindowManagerP
         dc.getDisplayPolicy().removeWindowLw(this);
 
         disposeInputChannel();
-        mSystemOnBackInvokedCallback = null;
+        mOnBackInvokedCallbackInfo = null;
-        mApplicationOnBackInvokedCallback = null;
 
         mSession.windowRemovedLocked();
         try {
@@ -2551,8 +2534,7 @@ class WindowState extends WindowContainer<WindowState> implements WindowManagerP
 
         try {
             disposeInputChannel();
-            mSystemOnBackInvokedCallback = null;
+            mOnBackInvokedCallbackInfo = null;
-            mApplicationOnBackInvokedCallback = null;
 
             ProtoLog.v(WM_DEBUG_APP_TRANSITIONS,
                     "Remove %s: mSurfaceController=%s mAnimatingExit=%b mRemoveOnExit=%b "

services/tests/wmtests/src/com/android/server/wm/BackNavigationControllerTests.java

@@ -44,6 +44,7 @@ import android.window.BackEvent;
 import android.window.BackNavigationInfo;
 import android.window.IOnBackInvokedCallback;
 import android.window.OnBackInvokedCallback;
+import android.window.OnBackInvokedCallbackInfo;
 import android.window.OnBackInvokedDispatcher;
 import android.window.TaskSnapshot;
 import android.window.WindowOnBackInvokedDispatcher;
@@ -104,7 +105,8 @@ public class BackNavigationControllerTests extends WindowTestsBase {
         WindowState window = createAppWindow(task, FIRST_APPLICATION_WINDOW, "window");
         addToWindowMap(window, true);
         IOnBackInvokedCallback callback = createOnBackInvokedCallback();
-        window.setOnBackInvokedCallback(callback, OnBackInvokedDispatcher.PRIORITY_SYSTEM);
+        window.setOnBackInvokedCallbackInfo(
+                new OnBackInvokedCallbackInfo(callback, OnBackInvokedDispatcher.PRIORITY_SYSTEM));
         BackNavigationInfo backNavigationInfo = startBackNavigation();
         assertWithMessage("BackNavigationInfo").that(backNavigationInfo).isNotNull();
         assertThat(typeToString(backNavigationInfo.getType()))
@@ -130,7 +132,8 @@ public class BackNavigationControllerTests extends WindowTestsBase {
         addToWindowMap(window, true);
 
         IOnBackInvokedCallback callback = createOnBackInvokedCallback();
-        window.setOnBackInvokedCallback(callback, OnBackInvokedDispatcher.PRIORITY_DEFAULT);
+        window.setOnBackInvokedCallbackInfo(
+                new OnBackInvokedCallbackInfo(callback, OnBackInvokedDispatcher.PRIORITY_DEFAULT));
 
         BackNavigationInfo backNavigationInfo = startBackNavigation();
         assertWithMessage("BackNavigationInfo").that(backNavigationInfo).isNotNull();
@@ -169,11 +172,9 @@ public class BackNavigationControllerTests extends WindowTestsBase {
         WindowState appWindow = task.getTopVisibleAppMainWindow();
         WindowOnBackInvokedDispatcher dispatcher = new WindowOnBackInvokedDispatcher();
         doAnswer(invocation -> {
-            appWindow.setOnBackInvokedCallback(invocation.getArgument(1),
+            appWindow.setOnBackInvokedCallbackInfo(invocation.getArgument(1));
-                    invocation.getArgument(2));
             return null;
-        }).when(appWindow.mSession).setOnBackInvokedCallback(eq(appWindow.mClient), any(),
+        }).when(appWindow.mSession).setOnBackInvokedCallbackInfo(eq(appWindow.mClient), any());
-                anyInt());
 
         addToWindowMap(appWindow, true);
         dispatcher.attachToWindow(appWindow.mSession, appWindow.mClient);
@@ -216,15 +217,15 @@ public class BackNavigationControllerTests extends WindowTestsBase {
 
     private IOnBackInvokedCallback withSystemCallback(Task task) {
         IOnBackInvokedCallback callback = createOnBackInvokedCallback();
-        task.getTopMostActivity().getTopChild().setOnBackInvokedCallback(callback,
+        task.getTopMostActivity().getTopChild().setOnBackInvokedCallbackInfo(
-                OnBackInvokedDispatcher.PRIORITY_SYSTEM);
+                new OnBackInvokedCallbackInfo(callback, OnBackInvokedDispatcher.PRIORITY_SYSTEM));
         return callback;
     }
 
     private IOnBackInvokedCallback withAppCallback(Task task) {
         IOnBackInvokedCallback callback = createOnBackInvokedCallback();
-        task.getTopMostActivity().getTopChild().setOnBackInvokedCallback(callback,
+        task.getTopMostActivity().getTopChild().setOnBackInvokedCallbackInfo(
-                OnBackInvokedDispatcher.PRIORITY_DEFAULT);
+                new OnBackInvokedCallbackInfo(callback, OnBackInvokedDispatcher.PRIORITY_DEFAULT));
         return callback;
     }