diff --git a/packages/SystemUI/shared/src/com/android/systemui/shared/system/InputConsumerController.java b/packages/SystemUI/shared/src/com/android/systemui/shared/system/InputConsumerController.java
index e554dcd..ebed1fc 100644
--- a/packages/SystemUI/shared/src/com/android/systemui/shared/system/InputConsumerController.java
+++ b/packages/SystemUI/shared/src/com/android/systemui/shared/system/InputConsumerController.java
@@ -64,7 +64,7 @@
     private final class InputEventReceiver extends BatchedInputEventReceiver {
 
         public InputEventReceiver(InputChannel inputChannel, Looper looper) {
-            super(inputChannel, looper, Choreographer.getSfInstance());
+            super(inputChannel, looper, Choreographer.getInstance());
         }
 
         @Override
diff --git a/packages/SystemUI/src/com/android/systemui/pip/PipAnimationController.java b/packages/SystemUI/src/com/android/systemui/pip/PipAnimationController.java
index 1fc1fe4..36b5fad 100644
--- a/packages/SystemUI/src/com/android/systemui/pip/PipAnimationController.java
+++ b/packages/SystemUI/src/com/android/systemui/pip/PipAnimationController.java
@@ -19,11 +19,8 @@
 import android.animation.Animator;
 import android.animation.ValueAnimator;
 import android.annotation.IntDef;
-import android.annotation.MainThread;
 import android.content.Context;
 import android.graphics.Rect;
-import android.os.RemoteException;
-import android.view.IWindowContainer;
 import android.view.SurfaceControl;
 import android.view.animation.AnimationUtils;
 import android.view.animation.Interpolator;
@@ -61,31 +58,30 @@
                 com.android.internal.R.interpolator.fast_out_slow_in);
     }
 
-    @MainThread
-    PipTransitionAnimator getAnimator(IWindowContainer wc, boolean scheduleFinishPip,
+    PipTransitionAnimator getAnimator(SurfaceControl leash, boolean scheduleFinishPip,
             Rect destinationBounds, float alphaStart, float alphaEnd) {
         if (mCurrentAnimator == null) {
             mCurrentAnimator = setupPipTransitionAnimator(
-                    PipTransitionAnimator.ofAlpha(wc, scheduleFinishPip,
-                            destinationBounds, alphaStart, alphaEnd));
+                    PipTransitionAnimator.ofAlpha(leash, scheduleFinishPip, destinationBounds,
+                            alphaStart, alphaEnd));
         } else if (mCurrentAnimator.getAnimationType() == ANIM_TYPE_ALPHA
                 && mCurrentAnimator.isRunning()) {
             mCurrentAnimator.updateEndValue(alphaEnd);
         } else {
             mCurrentAnimator.cancel();
             mCurrentAnimator = setupPipTransitionAnimator(
-                    PipTransitionAnimator.ofAlpha(wc, scheduleFinishPip,
-                            destinationBounds, alphaStart, alphaEnd));
+                    PipTransitionAnimator.ofAlpha(leash, scheduleFinishPip, destinationBounds,
+                            alphaStart, alphaEnd));
         }
         return mCurrentAnimator;
     }
 
-    @MainThread
-    PipTransitionAnimator getAnimator(IWindowContainer wc, boolean scheduleFinishPip,
+    PipTransitionAnimator getAnimator(SurfaceControl leash, boolean scheduleFinishPip,
             Rect startBounds, Rect endBounds) {
         if (mCurrentAnimator == null) {
             mCurrentAnimator = setupPipTransitionAnimator(
-                    PipTransitionAnimator.ofBounds(wc, scheduleFinishPip, startBounds, endBounds));
+                    PipTransitionAnimator.ofBounds(leash, scheduleFinishPip,
+                            startBounds, endBounds));
         } else if (mCurrentAnimator.getAnimationType() == ANIM_TYPE_BOUNDS
                 && mCurrentAnimator.isRunning()) {
             mCurrentAnimator.setDestinationBounds(endBounds);
@@ -94,7 +90,8 @@
         } else {
             mCurrentAnimator.cancel();
             mCurrentAnimator = setupPipTransitionAnimator(
-                    PipTransitionAnimator.ofBounds(wc, scheduleFinishPip, startBounds, endBounds));
+                    PipTransitionAnimator.ofBounds(leash, scheduleFinishPip,
+                            startBounds, endBounds));
         }
         return mCurrentAnimator;
     }
@@ -116,18 +113,18 @@
         /**
          * Called when PiP animation is started.
          */
-        public void onPipAnimationStart(IWindowContainer wc, PipTransitionAnimator animator) {}
+        public void onPipAnimationStart(PipTransitionAnimator animator) {}
 
         /**
          * Called when PiP animation is ended.
          */
-        public void onPipAnimationEnd(IWindowContainer wc, SurfaceControl.Transaction tx,
+        public void onPipAnimationEnd(SurfaceControl.Transaction tx,
                 PipTransitionAnimator animator) {}
 
         /**
          * Called when PiP animation is cancelled.
          */
-        public void onPipAnimationCancel(IWindowContainer wc, PipTransitionAnimator animator) {}
+        public void onPipAnimationCancel(PipTransitionAnimator animator) {}
     }
 
     /**
@@ -137,7 +134,6 @@
     public abstract static class PipTransitionAnimator<T> extends ValueAnimator implements
             ValueAnimator.AnimatorUpdateListener,
             ValueAnimator.AnimatorListener {
-        private final IWindowContainer mWindowContainer;
         private final boolean mScheduleFinishPip;
         private final SurfaceControl mLeash;
         private final @AnimationType int mAnimationType;
@@ -149,23 +145,18 @@
         private PipAnimationCallback mPipAnimationCallback;
         private SurfaceControlTransactionFactory mSurfaceControlTransactionFactory;
 
-        private PipTransitionAnimator(IWindowContainer wc, boolean scheduleFinishPip,
+        private PipTransitionAnimator(SurfaceControl leash, boolean scheduleFinishPip,
                 @AnimationType int animationType, Rect destinationBounds,
                 T startValue, T endValue) {
-            mWindowContainer = wc;
             mScheduleFinishPip = scheduleFinishPip;
-            try {
-                mLeash = wc.getLeash();
-                mAnimationType = animationType;
-                mDestinationBounds.set(destinationBounds);
-                mStartValue = startValue;
-                mEndValue = endValue;
-                addListener(this);
-                addUpdateListener(this);
-                mSurfaceControlTransactionFactory = SurfaceControl.Transaction::new;
-            } catch (RemoteException e) {
-                throw new RuntimeException(e);
-            }
+            mLeash = leash;
+            mAnimationType = animationType;
+            mDestinationBounds.set(destinationBounds);
+            mStartValue = startValue;
+            mEndValue = endValue;
+            addListener(this);
+            addUpdateListener(this);
+            mSurfaceControlTransactionFactory = SurfaceControl.Transaction::new;
         }
 
         @Override
@@ -173,7 +164,7 @@
             mCurrentValue = mStartValue;
             applySurfaceControlTransaction(mLeash, newSurfaceControlTransaction(), FRACTION_START);
             if (mPipAnimationCallback != null) {
-                mPipAnimationCallback.onPipAnimationStart(mWindowContainer, this);
+                mPipAnimationCallback.onPipAnimationStart(this);
             }
         }
 
@@ -189,14 +180,14 @@
             final SurfaceControl.Transaction tx = newSurfaceControlTransaction();
             applySurfaceControlTransaction(mLeash, tx, FRACTION_END);
             if (mPipAnimationCallback != null) {
-                mPipAnimationCallback.onPipAnimationEnd(mWindowContainer, tx, this);
+                mPipAnimationCallback.onPipAnimationEnd(tx, this);
             }
         }
 
         @Override
         public void onAnimationCancel(Animator animation) {
             if (mPipAnimationCallback != null) {
-                mPipAnimationCallback.onPipAnimationCancel(mWindowContainer, this);
+                mPipAnimationCallback.onPipAnimationCancel(this);
             }
         }
 
@@ -260,9 +251,9 @@
         abstract void applySurfaceControlTransaction(SurfaceControl leash,
                 SurfaceControl.Transaction tx, float fraction);
 
-        static PipTransitionAnimator<Float> ofAlpha(IWindowContainer wc, boolean scheduleFinishPip,
+        static PipTransitionAnimator<Float> ofAlpha(SurfaceControl leash, boolean scheduleFinishPip,
                 Rect destinationBounds, float startValue, float endValue) {
-            return new PipTransitionAnimator<Float>(wc, scheduleFinishPip, ANIM_TYPE_ALPHA,
+            return new PipTransitionAnimator<Float>(leash, scheduleFinishPip, ANIM_TYPE_ALPHA,
                     destinationBounds, startValue, endValue) {
                 @Override
                 void applySurfaceControlTransaction(SurfaceControl leash,
@@ -281,10 +272,10 @@
             };
         }
 
-        static PipTransitionAnimator<Rect> ofBounds(IWindowContainer wc, boolean scheduleFinishPip,
+        static PipTransitionAnimator<Rect> ofBounds(SurfaceControl leash, boolean scheduleFinishPip,
                 Rect startValue, Rect endValue) {
             // construct new Rect instances in case they are recycled
-            return new PipTransitionAnimator<Rect>(wc, scheduleFinishPip, ANIM_TYPE_BOUNDS,
+            return new PipTransitionAnimator<Rect>(leash, scheduleFinishPip, ANIM_TYPE_BOUNDS,
                     endValue, new Rect(startValue), new Rect(endValue)) {
                 private final Rect mTmpRect = new Rect();
 
diff --git a/packages/SystemUI/src/com/android/systemui/pip/PipTaskOrganizer.java b/packages/SystemUI/src/com/android/systemui/pip/PipTaskOrganizer.java
index 836485a..4766ebc 100644
--- a/packages/SystemUI/src/com/android/systemui/pip/PipTaskOrganizer.java
+++ b/packages/SystemUI/src/com/android/systemui/pip/PipTaskOrganizer.java
@@ -30,6 +30,7 @@
 import android.graphics.Rect;
 import android.os.Handler;
 import android.os.Looper;
+import android.os.Message;
 import android.os.RemoteException;
 import android.util.Log;
 import android.view.DisplayInfo;
@@ -38,9 +39,13 @@
 import android.view.SurfaceControl;
 import android.view.WindowContainerTransaction;
 
+import com.android.internal.os.SomeArgs;
+import com.android.systemui.pip.phone.PipUpdateThread;
+
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
+import java.util.function.Consumer;
 
 /**
  * Manages PiP tasks such as resize and offset.
@@ -56,7 +61,13 @@
 public class PipTaskOrganizer extends ITaskOrganizer.Stub {
     private static final String TAG = PipTaskOrganizer.class.getSimpleName();
 
+    private static final int MSG_RESIZE_IMMEDIATE = 1;
+    private static final int MSG_RESIZE_ANIMATE = 2;
+    private static final int MSG_OFFSET_ANIMATE = 3;
+    private static final int MSG_FINISH_RESIZE = 4;
+
     private final Handler mMainHandler;
+    private final Handler mUpdateHandler;
     private final ITaskOrganizerController mTaskOrganizerController;
     private final PipBoundsHandler mPipBoundsHandler;
     private final PipAnimationController mPipAnimationController;
@@ -64,11 +75,11 @@
     private final Rect mDisplayBounds = new Rect();
     private final Rect mLastReportedBounds = new Rect();
 
+    // These callbacks are called on the update thread
     private final PipAnimationController.PipAnimationCallback mPipAnimationCallback =
             new PipAnimationController.PipAnimationCallback() {
         @Override
-        public void onPipAnimationStart(IWindowContainer wc,
-                PipAnimationController.PipTransitionAnimator animator) {
+        public void onPipAnimationStart(PipAnimationController.PipTransitionAnimator animator) {
             mMainHandler.post(() -> {
                 for (int i = mPipTransitionCallbacks.size() - 1; i >= 0; i--) {
                     final PipTransitionCallback callback = mPipTransitionCallbacks.get(i);
@@ -78,7 +89,7 @@
         }
 
         @Override
-        public void onPipAnimationEnd(IWindowContainer wc, SurfaceControl.Transaction tx,
+        public void onPipAnimationEnd(SurfaceControl.Transaction tx,
                 PipAnimationController.PipTransitionAnimator animator) {
             mMainHandler.post(() -> {
                 for (int i = mPipTransitionCallbacks.size() - 1; i >= 0; i--) {
@@ -86,13 +97,11 @@
                     callback.onPipTransitionFinished();
                 }
             });
-            final Rect destinationBounds = animator.getDestinationBounds();
-            finishResizeInternal(destinationBounds, wc, tx, animator.shouldScheduleFinishPip());
+            finishResize(animator.getDestinationBounds(), tx, animator.shouldScheduleFinishPip());
         }
 
         @Override
-        public void onPipAnimationCancel(IWindowContainer wc,
-                PipAnimationController.PipTransitionAnimator animator) {
+        public void onPipAnimationCancel(PipAnimationController.PipTransitionAnimator animator) {
             mMainHandler.post(() -> {
                 for (int i = mPipTransitionCallbacks.size() - 1; i >= 0; i--) {
                     final PipTransitionCallback callback = mPipTransitionCallbacks.get(i);
@@ -102,28 +111,75 @@
         }
     };
 
+    private Handler.Callback mUpdateCallbacks = new Handler.Callback() {
+        @Override
+        public boolean handleMessage(Message msg) {
+            SomeArgs args = (SomeArgs) msg.obj;
+            Consumer<Rect> updateBoundsCallback = (Consumer<Rect>) args.arg1;
+            switch (msg.what) {
+                case MSG_RESIZE_IMMEDIATE: {
+                    Rect toBounds = (Rect) args.arg2;
+                    resizePip(toBounds);
+                    if (updateBoundsCallback != null) {
+                        updateBoundsCallback.accept(toBounds);
+                    }
+                    break;
+                }
+                case MSG_RESIZE_ANIMATE: {
+                    Rect currentBounds = (Rect) args.arg2;
+                    Rect toBounds = (Rect) args.arg3;
+                    boolean scheduleFinishPip = args.argi1 != 0;
+                    int duration = args.argi2;
+                    animateResizePip(scheduleFinishPip, currentBounds, toBounds, duration);
+                    if (updateBoundsCallback != null) {
+                        updateBoundsCallback.accept(toBounds);
+                    }
+                    break;
+                }
+                case MSG_OFFSET_ANIMATE: {
+                    Rect originalBounds = (Rect) args.arg2;
+                    final int offset = args.argi1;
+                    final int duration = args.argi2;
+                    offsetPip(originalBounds, 0 /* xOffset */, offset, duration);
+                    Rect toBounds = new Rect(originalBounds);
+                    toBounds.offset(0, offset);
+                    if (updateBoundsCallback != null) {
+                        updateBoundsCallback.accept(toBounds);
+                    }
+                    break;
+                }
+                case MSG_FINISH_RESIZE: {
+                    SurfaceControl.Transaction tx = (SurfaceControl.Transaction) args.arg2;
+                    Rect toBounds = (Rect) args.arg3;
+                    boolean scheduleFinishPip = args.argi1 != 0;
+                    finishResize(toBounds, tx, scheduleFinishPip);
+                    if (updateBoundsCallback != null) {
+                        updateBoundsCallback.accept(toBounds);
+                    }
+                    break;
+                }
+            }
+            args.recycle();
+            return true;
+        }
+    };
+
     private ActivityManager.RunningTaskInfo mTaskInfo;
+    private IWindowContainer mToken;
+    private SurfaceControl mLeash;
+    private boolean mInPip;
     private @PipAnimationController.AnimationType int mOneShotAnimationType = ANIM_TYPE_BOUNDS;
 
     public PipTaskOrganizer(Context context, @NonNull PipBoundsHandler boundsHandler) {
         mMainHandler = new Handler(Looper.getMainLooper());
+        mUpdateHandler = new Handler(PipUpdateThread.get().getLooper(), mUpdateCallbacks);
         mTaskOrganizerController = ActivityTaskManager.getTaskOrganizerController();
         mPipBoundsHandler = boundsHandler;
         mPipAnimationController = new PipAnimationController(context);
     }
 
-    /**
-     * Offset the PiP window, animate if the given duration is not {@link #DURATION_NONE}
-     */
-    public void offsetPinnedStack(Rect originalBounds, int xOffset, int yOffset, int durationMs) {
-        if (mTaskInfo == null) {
-            Log.w(TAG, "mTaskInfo is not set");
-            return;
-        }
-        final Rect destinationBounds = new Rect(originalBounds);
-        destinationBounds.offset(xOffset, yOffset);
-        animateResizePipInternal(mTaskInfo.token, false /* scheduleFinishPip*/,
-                originalBounds, destinationBounds, durationMs);
+    public Handler getUpdateHandler() {
+        return mUpdateHandler;
     }
 
     /**
@@ -171,7 +227,7 @@
         try {
             mLastReportedBounds.set(destinationBounds);
             final WindowContainerTransaction wct = new WindowContainerTransaction();
-            wct.setBounds(mTaskInfo.token, destinationBounds);
+            wct.setBounds(mToken, destinationBounds);
             mTaskOrganizerController.applyContainerTransaction(wct, null /* ITaskOrganizer */);
         } catch (RemoteException e) {
             Log.w(TAG, "Failed to apply window container transaction", e);
@@ -185,13 +241,20 @@
                 getAspectRatioOrDefault(info.pictureInPictureParams), null /* bounds */);
         Objects.requireNonNull(destinationBounds, "Missing destination bounds");
         mTaskInfo = info;
+        mToken = mTaskInfo.token;
+        mInPip = true;
+        try {
+            mLeash = mToken.getLeash();
+        } catch (RemoteException e) {
+            throw new RuntimeException("Unable to get leash", e);
+        }
         if (mOneShotAnimationType == ANIM_TYPE_BOUNDS) {
             final Rect currentBounds = mTaskInfo.configuration.windowConfiguration.getBounds();
-            animateResizePipInternal(mTaskInfo.token, true /* scheduleFinishPip */,
-                    currentBounds, destinationBounds, DURATION_DEFAULT_MS);
+            scheduleAnimateResizePip(true /* scheduleFinishPip */,
+                    currentBounds, destinationBounds, DURATION_DEFAULT_MS, null);
         } else if (mOneShotAnimationType == ANIM_TYPE_ALPHA) {
-            mMainHandler.post(() -> mPipAnimationController
-                    .getAnimator(mTaskInfo.token, true /* scheduleFinishPip */,
+            mUpdateHandler.post(() -> mPipAnimationController
+                    .getAnimator(mLeash, true /* scheduleFinishPip */,
                             destinationBounds, 0f, 1f)
                     .setPipAnimationCallback(mPipAnimationCallback)
                     .setDuration(DURATION_DEFAULT_MS)
@@ -205,12 +268,12 @@
     @Override
     public void taskVanished(IWindowContainer token) {
         Objects.requireNonNull(token, "Requires valid IWindowContainer");
-        if (token.asBinder() != mTaskInfo.token.asBinder()) {
+        if (token.asBinder() != mToken.asBinder()) {
             Log.wtf(TAG, "Unrecognized token: " + token);
             return;
         }
-        animateResizePipInternal(token, false /* scheduleFinishPip */,
-                mLastReportedBounds, mDisplayBounds, DURATION_DEFAULT_MS);
+        scheduleAnimateResizePip(mDisplayBounds, DURATION_DEFAULT_MS, null);
+        mInPip = false;
     }
 
     @Override
@@ -227,7 +290,7 @@
         final Rect destinationBounds = mPipBoundsHandler.getDestinationBounds(
                 getAspectRatioOrDefault(newParams), null /* bounds */);
         Objects.requireNonNull(destinationBounds, "Missing destination bounds");
-        animateResizePip(destinationBounds, DURATION_DEFAULT_MS);
+        scheduleAnimateResizePip(destinationBounds, DURATION_DEFAULT_MS, null);
     }
 
     /**
@@ -243,102 +306,158 @@
     }
 
     /**
-     * Directly perform manipulation/resize on the leash. This will not perform any
-     * {@link WindowContainerTransaction} until {@link #finishResize} is called.
+     * Animates resizing of the pinned stack given the duration.
      */
-    public void resizePip(Rect destinationBounds) {
-        Objects.requireNonNull(mTaskInfo, "Requires valid IWindowContainer");
-        resizePipInternal(mTaskInfo.token, destinationBounds);
+    public void scheduleAnimateResizePip(Rect toBounds, int duration,
+            Consumer<Rect> updateBoundsCallback) {
+        scheduleAnimateResizePip(false /* scheduleFinishPip */,
+                mLastReportedBounds, toBounds, duration, updateBoundsCallback);
     }
 
-    private void resizePipInternal(IWindowContainer wc,
-            Rect destinationBounds) {
-        Objects.requireNonNull(mTaskInfo, "Requires valid IWindowContainer");
-        try {
-            // Could happen when dismissPip
-            if (wc == null || wc.getLeash() == null) {
-                Log.w(TAG, "Abort animation, invalid leash");
-                return;
-            }
-            final SurfaceControl leash = wc.getLeash();
-            new SurfaceControl.Transaction()
-                    .setPosition(leash, destinationBounds.left, destinationBounds.top)
-                    .setWindowCrop(leash, destinationBounds.width(), destinationBounds.height())
-                    .apply();
-        } catch (RemoteException e) {
-            Log.w(TAG, "Abort animation, invalid window container", e);
-        } catch (Exception e) {
-            Log.e(TAG, "Should not reach here, terrible thing happened", e);
+    private void scheduleAnimateResizePip(boolean scheduleFinishPip,
+            Rect currentBounds, Rect destinationBounds, int durationMs,
+            Consumer<Rect> updateBoundsCallback) {
+        Objects.requireNonNull(mToken, "Requires valid IWindowContainer");
+        if (!mInPip) {
+            // Ignore animation when we are no longer in PIP
+            return;
         }
+        SomeArgs args = SomeArgs.obtain();
+        args.arg1 = updateBoundsCallback;
+        args.arg2 = currentBounds;
+        args.arg3 = destinationBounds;
+        args.argi1 = scheduleFinishPip ? 1 : 0;
+        args.argi2 = durationMs;
+        mUpdateHandler.sendMessage(mUpdateHandler.obtainMessage(MSG_RESIZE_ANIMATE, args));
+    }
+
+    /**
+     * Directly perform manipulation/resize on the leash. This will not perform any
+     * {@link WindowContainerTransaction} until {@link #scheduleFinishResizePip} is called.
+     */
+    public void scheduleResizePip(Rect toBounds, Consumer<Rect> updateBoundsCallback) {
+        Objects.requireNonNull(mToken, "Requires valid IWindowContainer");
+        SomeArgs args = SomeArgs.obtain();
+        args.arg1 = updateBoundsCallback;
+        args.arg2 = toBounds;
+        mUpdateHandler.sendMessage(mUpdateHandler.obtainMessage(MSG_RESIZE_IMMEDIATE, args));
     }
 
     /**
      * Finish a intermediate resize operation. This is expected to be called after
-     * {@link #resizePip}.
+     * {@link #scheduleResizePip}.
      */
-    public void finishResize(Rect destinationBounds) {
-        try {
-            final IWindowContainer wc = mTaskInfo.token;
-            SurfaceControl.Transaction tx = new SurfaceControl.Transaction()
-                    .setPosition(wc.getLeash(), destinationBounds.left,
-                            destinationBounds.top)
-                    .setWindowCrop(wc.getLeash(), destinationBounds.width(),
-                            destinationBounds.height());
-            finishResizeInternal(destinationBounds, wc, tx, false);
-        } catch (RemoteException e) {
-            Log.e(TAG, "Failed to obtain leash");
-        }
+    public void scheduleFinishResizePip(Rect destinationBounds) {
+        Objects.requireNonNull(mToken, "Requires valid IWindowContainer");
+        SurfaceControl.Transaction tx = new SurfaceControl.Transaction()
+                .setPosition(mLeash, destinationBounds.left, destinationBounds.top)
+                .setWindowCrop(mLeash, destinationBounds.width(), destinationBounds.height());
+        scheduleFinishResizePip(tx, destinationBounds, false /* scheduleFinishPip */,
+                null);
     }
 
-    private void finishResizeInternal(Rect destinationBounds, IWindowContainer wc,
-            SurfaceControl.Transaction tx, boolean shouldScheduleFinishPip) {
+    private void scheduleFinishResizePip(SurfaceControl.Transaction tx,
+            Rect destinationBounds, boolean scheduleFinishPip,
+            Consumer<Rect> updateBoundsCallback) {
+        Objects.requireNonNull(mToken, "Requires valid IWindowContainer");
+        SomeArgs args = SomeArgs.obtain();
+        args.arg1 = updateBoundsCallback;
+        args.arg2 = tx;
+        args.arg3 = destinationBounds;
+        args.argi1 = scheduleFinishPip ? 1 : 0;
+        mUpdateHandler.sendMessage(mUpdateHandler.obtainMessage(MSG_FINISH_RESIZE, args));
+    }
+
+    /**
+     * Offset the PiP window, animate if the given duration is not {@link #DURATION_NONE}
+     */
+    public void scheduleOffsetPip(Rect originalBounds, int offset, int duration,
+            Consumer<Rect> updateBoundsCallback) {
+        if (!mInPip) {
+            // Ignore offsets when we are no longer in PIP
+            return;
+        }
+        Objects.requireNonNull(mToken, "Requires valid IWindowContainer");
+        SomeArgs args = SomeArgs.obtain();
+        args.arg1 = updateBoundsCallback;
+        args.arg2 = originalBounds;
+        // offset would be zero if triggered from screen rotation.
+        args.argi1 = offset;
+        args.argi2 = duration;
+        mUpdateHandler.sendMessage(mUpdateHandler.obtainMessage(MSG_OFFSET_ANIMATE, args));
+    }
+
+    private void offsetPip(Rect originalBounds, int xOffset, int yOffset, int durationMs) {
+        if (Looper.myLooper() != mUpdateHandler.getLooper()) {
+            throw new RuntimeException("Callers should call scheduleOffsetPip() instead of this "
+                    + "directly");
+        }
+        if (mTaskInfo == null) {
+            Log.w(TAG, "mTaskInfo is not set");
+            return;
+        }
+        final Rect destinationBounds = new Rect(originalBounds);
+        destinationBounds.offset(xOffset, yOffset);
+        animateResizePip(false /* scheduleFinishPip*/, originalBounds, destinationBounds,
+                durationMs);
+    }
+
+    private void resizePip(Rect destinationBounds) {
+        if (Looper.myLooper() != mUpdateHandler.getLooper()) {
+            throw new RuntimeException("Callers should call scheduleResizePip() instead of this "
+                    + "directly");
+        }
+        Objects.requireNonNull(mToken, "Requires valid IWindowContainer");
+        // Could happen when dismissPip
+        if (mToken == null || mLeash == null) {
+            Log.w(TAG, "Abort animation, invalid leash");
+            return;
+        }
+        new SurfaceControl.Transaction()
+                .setPosition(mLeash, destinationBounds.left, destinationBounds.top)
+                .setWindowCrop(mLeash, destinationBounds.width(), destinationBounds.height())
+                .apply();
+    }
+
+    private void finishResize(Rect destinationBounds, SurfaceControl.Transaction tx,
+            boolean shouldScheduleFinishPip) {
+        if (Looper.myLooper() != mUpdateHandler.getLooper()) {
+            throw new RuntimeException("Callers should call scheduleResizePip() instead of this "
+                    + "directly");
+        }
         mLastReportedBounds.set(destinationBounds);
         try {
             final WindowContainerTransaction wct = new WindowContainerTransaction();
             if (shouldScheduleFinishPip) {
-                wct.scheduleFinishEnterPip(wc, destinationBounds);
+                wct.scheduleFinishEnterPip(mToken, destinationBounds);
             } else {
-                wct.setBounds(wc, destinationBounds);
+                wct.setBounds(mToken, destinationBounds);
             }
-            wct.setBoundsChangeTransaction(mTaskInfo.token, tx);
+            wct.setBoundsChangeTransaction(mToken, tx);
             mTaskOrganizerController.applyContainerTransaction(wct, null /* ITaskOrganizer */);
         } catch (RemoteException e) {
             Log.e(TAG, "Failed to apply container transaction", e);
         }
     }
 
-    /**
-     * Animates resizing of the pinned stack given the duration.
-     */
-    public void animateResizePip(Rect destinationBounds, int durationMs) {
-        Objects.requireNonNull(mTaskInfo, "Requires valid IWindowContainer");
-        animateResizePipInternal(mTaskInfo.token, false, mLastReportedBounds,
-                destinationBounds, durationMs);
-    }
-
-    private void animateResizePipInternal(IWindowContainer wc, boolean scheduleFinishPip,
-            Rect currentBounds, Rect destinationBounds, int durationMs) {
-        try {
-            // Could happen when dismissPip
-            if (wc == null || wc.getLeash() == null) {
-                Log.w(TAG, "Abort animation, invalid leash");
-                return;
-            }
-            final SurfaceControl leash = wc.getLeash();
-
-            mMainHandler.post(() -> mPipAnimationController
-                    .getAnimator(wc, scheduleFinishPip, currentBounds, destinationBounds)
-                    .setPipAnimationCallback(mPipAnimationCallback)
-                    .setDuration(durationMs)
-                    .start());
-        } catch (RemoteException e) {
-            Log.w(TAG, "Abort animation, invalid window container", e);
-        } catch (Exception e) {
-            Log.e(TAG, "Should not reach here, terrible thing happened", e);
+    private void animateResizePip(boolean scheduleFinishPip, Rect currentBounds,
+            Rect destinationBounds, int durationMs) {
+        if (Looper.myLooper() != mUpdateHandler.getLooper()) {
+            throw new RuntimeException("Callers should call scheduleAnimateResizePip() instead of "
+                    + "this directly");
         }
+        // Could happen when dismissPip
+        if (mToken == null || mLeash == null) {
+            Log.w(TAG, "Abort animation, invalid leash");
+            return;
+        }
+        mUpdateHandler.post(() -> mPipAnimationController
+                .getAnimator(mLeash, scheduleFinishPip, currentBounds, destinationBounds)
+                .setPipAnimationCallback(mPipAnimationCallback)
+                .setDuration(durationMs)
+                .start());
     }
 
-
     private float getAspectRatioOrDefault(@Nullable PictureInPictureParams params) {
         return params == null
                 ? mPipBoundsHandler.getDefaultAspectRatio()
diff --git a/packages/SystemUI/src/com/android/systemui/pip/phone/PipMotionHelper.java b/packages/SystemUI/src/com/android/systemui/pip/phone/PipMotionHelper.java
index fdaf66a..33760be 100644
--- a/packages/SystemUI/src/com/android/systemui/pip/phone/PipMotionHelper.java
+++ b/packages/SystemUI/src/com/android/systemui/pip/phone/PipMotionHelper.java
@@ -28,16 +28,11 @@
 import android.graphics.PointF;
 import android.graphics.Rect;
 import android.os.Debug;
-import android.os.Handler;
-import android.os.Message;
 import android.os.RemoteException;
 import android.util.Log;
-import android.view.Choreographer;
 
 import androidx.dynamicanimation.animation.SpringForce;
 
-import com.android.internal.graphics.SfVsyncFrameCallbackProvider;
-import com.android.internal.os.SomeArgs;
 import com.android.systemui.pip.PipSnapAlgorithm;
 import com.android.systemui.pip.PipTaskOrganizer;
 import com.android.systemui.shared.system.WindowManagerWrapper;
@@ -47,11 +42,12 @@
 import com.android.systemui.util.animation.PhysicsAnimator;
 
 import java.io.PrintWriter;
+import java.util.function.Consumer;
 
 /**
  * A helper to animate and manipulate the PiP.
  */
-public class PipMotionHelper implements Handler.Callback, PipAppOpsListener.Callback,
+public class PipMotionHelper implements PipAppOpsListener.Callback,
         FloatingContentCoordinator.FloatingContent {
 
     private static final String TAG = "PipMotionHelper";
@@ -68,14 +64,9 @@
     // The fraction of the stack height that the user has to drag offscreen to dismiss the PiP
     private static final float DISMISS_OFFSCREEN_FRACTION = 0.3f;
 
-    private static final int MSG_RESIZE_IMMEDIATE = 1;
-    private static final int MSG_RESIZE_ANIMATE = 2;
-    private static final int MSG_OFFSET_ANIMATE = 3;
-
     private final Context mContext;
     private final IActivityTaskManager mActivityTaskManager;
     private final PipTaskOrganizer mPipTaskOrganizer;
-    private final Handler mHandler;
 
     private PipMenuActivityController mMenuController;
     private PipSnapAlgorithm mSnapAlgorithm;
@@ -92,9 +83,6 @@
     /** The region that all of PIP must stay within. */
     private Rect mFloatingAllowedArea = new Rect();
 
-    private final SfVsyncFrameCallbackProvider mSfVsyncFrameProvider =
-            new SfVsyncFrameCallbackProvider();
-
     /**
      * Bounds that are animated using the physics animator.
      */
@@ -112,16 +100,11 @@
     private PhysicsAnimator<Rect> mAnimatedBoundsPhysicsAnimator = PhysicsAnimator.getInstance(
             mAnimatedBounds);
 
-    /** Callback that re-sizes PIP to the animated bounds. */
-    private final Choreographer.FrameCallback mResizePipVsyncCallback =
-            l -> resizePipUnchecked(mAnimatedBounds);
-
     /**
-     * Update listener that posts a vsync frame callback to resize PIP to {@link #mAnimatedBounds}.
+     * Update listener that resizes the PIP to {@link #mAnimatedBounds}.
      */
-    private final PhysicsAnimator.UpdateListener<Rect> mResizePipVsyncUpdateListener =
-            (target, values) ->
-                    mSfVsyncFrameProvider.postFrameCallback(mResizePipVsyncCallback);
+    private final PhysicsAnimator.UpdateListener<Rect> mResizePipUpdateListener =
+            (target, values) -> resizePipUnchecked(mAnimatedBounds);
 
     /** FlingConfig instances provided to PhysicsAnimator for fling gestures. */
     private PhysicsAnimator.FlingConfig mFlingConfigX;
@@ -137,12 +120,13 @@
                 new PhysicsAnimator.SpringConfig(
                         SpringForce.STIFFNESS_LOW, SpringForce.DAMPING_RATIO_LOW_BOUNCY);
 
+    private final Consumer<Rect> mUpdateBoundsCallback = (toBounds) -> mBounds.set(toBounds);
+
     public PipMotionHelper(Context context, IActivityTaskManager activityTaskManager,
             PipTaskOrganizer pipTaskOrganizer, PipMenuActivityController menuController,
             PipSnapAlgorithm snapAlgorithm, FlingAnimationUtils flingAnimationUtils,
             FloatingContentCoordinator floatingContentCoordinator) {
         mContext = context;
-        mHandler = new Handler(ForegroundThread.get().getLooper(), this);
         mActivityTaskManager = activityTaskManager;
         mPipTaskOrganizer = pipTaskOrganizer;
         mMenuController = menuController;
@@ -234,7 +218,7 @@
         }
         cancelAnimations();
         mMenuController.hideMenuWithoutResize();
-        mHandler.post(() -> {
+        mPipTaskOrganizer.getUpdateHandler().post(() -> {
             try {
                 mActivityTaskManager.dismissPip(!skipAnimation, EXPAND_STACK_TO_FULLSCREEN_DURATION);
             } catch (RemoteException e) {
@@ -253,7 +237,7 @@
         }
         cancelAnimations();
         mMenuController.hideMenuWithoutResize();
-        mHandler.post(() -> {
+        mPipTaskOrganizer.getUpdateHandler().post(() -> {
             try {
                 mActivityTaskManager.removeStacksInWindowingModes(
                         new int[]{ WINDOWING_MODE_PINNED });
@@ -406,17 +390,13 @@
      * Animates the PiP to offset it from the IME or shelf.
      */
     void animateToOffset(Rect originalBounds, int offset) {
+        if (DEBUG) {
+            Log.d(TAG, "animateToOffset: originalBounds=" + originalBounds + " offset=" + offset
+                    + " callers=\n" + Debug.getCallers(5, "    "));
+        }
         cancelAnimations();
-        adjustAndAnimatePipOffset(originalBounds, offset, SHIFT_DURATION);
-    }
-
-    private void adjustAndAnimatePipOffset(Rect originalBounds, int offset, int duration) {
-        SomeArgs args = SomeArgs.obtain();
-        args.arg1 = originalBounds;
-        // offset would be zero if triggered from screen rotation.
-        args.argi1 = offset;
-        args.argi2 = duration;
-        mHandler.sendMessage(mHandler.obtainMessage(MSG_OFFSET_ANIMATE, args));
+        mPipTaskOrganizer.scheduleOffsetPip(originalBounds, offset, SHIFT_DURATION,
+                mUpdateBoundsCallback);
     }
 
     /**
@@ -437,8 +417,7 @@
 
     /**
      * Starts the physics animator which will update the animated PIP bounds using physics
-     * animations, as well as the TimeAnimator which will apply those bounds to PIP at intervals
-     * synchronized with the SurfaceFlinger vsync frame provider.
+     * animations, as well as the TimeAnimator which will apply those bounds to PIP.
      *
      * This will also add end actions to the bounds animator that cancel the TimeAnimator and update
      * the 'real' bounds to equal the final animated bounds.
@@ -448,7 +427,7 @@
 
         mAnimatedBoundsPhysicsAnimator
                 .withEndActions(() ->  mPipTaskOrganizer.onMotionMovementEnd(mAnimatedBounds))
-                .addUpdateListener(mResizePipVsyncUpdateListener)
+                .addUpdateListener(mResizePipUpdateListener)
                 .start();
     }
 
@@ -471,9 +450,7 @@
                     + " callers=\n" + Debug.getCallers(5, "    "));
         }
         if (!toBounds.equals(mBounds)) {
-            SomeArgs args = SomeArgs.obtain();
-            args.arg1 = toBounds;
-            mHandler.sendMessage(mHandler.obtainMessage(MSG_RESIZE_IMMEDIATE, args));
+            mPipTaskOrganizer.scheduleResizePip(toBounds, mUpdateBoundsCallback);
         }
     }
 
@@ -486,10 +463,7 @@
                     + " duration=" + duration + " callers=\n" + Debug.getCallers(5, "    "));
         }
         if (!toBounds.equals(mBounds)) {
-            SomeArgs args = SomeArgs.obtain();
-            args.arg1 = toBounds;
-            args.argi1 = duration;
-            mHandler.sendMessage(mHandler.obtainMessage(MSG_RESIZE_ANIMATE, args));
+            mPipTaskOrganizer.scheduleAnimateResizePip(toBounds, duration, mUpdateBoundsCallback);
             setAnimatingToBounds(toBounds);
         }
     }
@@ -538,70 +512,6 @@
         return dismissArea.contains(endpoint.x, endpoint.y);
     }
 
-    /**
-     * Handles messages to be processed on the background thread.
-     */
-    public boolean handleMessage(Message msg) {
-        switch (msg.what) {
-            case MSG_RESIZE_IMMEDIATE: {
-                SomeArgs args = (SomeArgs) msg.obj;
-                Rect toBounds = (Rect) args.arg1;
-                mPipTaskOrganizer.resizePip(toBounds);
-                mBounds.set(toBounds);
-                return true;
-            }
-
-            case MSG_RESIZE_ANIMATE: {
-                SomeArgs args = (SomeArgs) msg.obj;
-                Rect toBounds = (Rect) args.arg1;
-                int duration = args.argi1;
-                try {
-                    StackInfo stackInfo = mActivityTaskManager.getStackInfo(
-                            WINDOWING_MODE_PINNED, ACTIVITY_TYPE_UNDEFINED);
-                    if (stackInfo == null) {
-                        // In the case where we've already re-expanded or dismissed the PiP, then
-                        // just skip the resize
-                        return true;
-                    }
-
-                    mPipTaskOrganizer.animateResizePip(toBounds, duration);
-                    mBounds.set(toBounds);
-                } catch (RemoteException e) {
-                    Log.e(TAG, "Could not animate resize pinned stack to bounds: " + toBounds, e);
-                }
-                return true;
-            }
-
-            case MSG_OFFSET_ANIMATE: {
-                SomeArgs args = (SomeArgs) msg.obj;
-                Rect originalBounds = (Rect) args.arg1;
-                final int offset = args.argi1;
-                final int duration = args.argi2;
-                try {
-                    StackInfo stackInfo = mActivityTaskManager.getStackInfo(
-                            WINDOWING_MODE_PINNED, ACTIVITY_TYPE_UNDEFINED);
-                    if (stackInfo == null) {
-                        // In the case where we've already re-expanded or dismissed the PiP, then
-                        // just skip the resize
-                        return true;
-                    }
-
-                    mPipTaskOrganizer.offsetPinnedStack(originalBounds,
-                            0 /* xOffset */, offset, duration);
-                    Rect toBounds = new Rect(originalBounds);
-                    toBounds.offset(0, offset);
-                    mBounds.set(toBounds);
-                } catch (RemoteException e) {
-                    Log.e(TAG, "Could not animate offset pinned stack with offset: " + offset, e);
-                }
-                return true;
-            }
-
-            default:
-                return false;
-        }
-    }
-
     public void dump(PrintWriter pw, String prefix) {
         final String innerPrefix = prefix + "  ";
         pw.println(prefix + TAG);
diff --git a/packages/SystemUI/src/com/android/systemui/pip/phone/ForegroundThread.java b/packages/SystemUI/src/com/android/systemui/pip/phone/PipUpdateThread.java
similarity index 70%
rename from packages/SystemUI/src/com/android/systemui/pip/phone/ForegroundThread.java
rename to packages/SystemUI/src/com/android/systemui/pip/phone/PipUpdateThread.java
index 9bf46bb..6c5d846 100644
--- a/packages/SystemUI/src/com/android/systemui/pip/phone/ForegroundThread.java
+++ b/packages/SystemUI/src/com/android/systemui/pip/phone/PipUpdateThread.java
@@ -21,33 +21,38 @@
 
 /**
  * Similar to {@link com.android.internal.os.BackgroundThread}, this is a shared singleton
- * foreground thread for each process.
+ * foreground thread for each process for updating PIP.
  */
-public final class ForegroundThread extends HandlerThread {
-    private static ForegroundThread sInstance;
+public final class PipUpdateThread extends HandlerThread {
+    private static PipUpdateThread sInstance;
     private static Handler sHandler;
 
-    private ForegroundThread() {
-        super("recents.fg");
+    private PipUpdateThread() {
+        super("pip");
     }
 
     private static void ensureThreadLocked() {
         if (sInstance == null) {
-            sInstance = new ForegroundThread();
+            sInstance = new PipUpdateThread();
             sInstance.start();
             sHandler = new Handler(sInstance.getLooper());
         }
     }
 
-    public static ForegroundThread get() {
-        synchronized (ForegroundThread.class) {
+    /**
+     * @return the static update thread instance
+     */
+    public static PipUpdateThread get() {
+        synchronized (PipUpdateThread.class) {
             ensureThreadLocked();
             return sInstance;
         }
     }
-
+    /**
+     * @return the static update thread handler instance
+     */
     public static Handler getHandler() {
-        synchronized (ForegroundThread.class) {
+        synchronized (PipUpdateThread.class) {
             ensureThreadLocked();
             return sHandler;
         }
diff --git a/packages/SystemUI/src/com/android/systemui/pip/tv/PipManager.java b/packages/SystemUI/src/com/android/systemui/pip/tv/PipManager.java
index f28c3f6..8636172 100644
--- a/packages/SystemUI/src/com/android/systemui/pip/tv/PipManager.java
+++ b/packages/SystemUI/src/com/android/systemui/pip/tv/PipManager.java
@@ -20,6 +20,8 @@
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_UNDEFINED;
 import static android.app.WindowConfiguration.WINDOWING_MODE_PINNED;
 
+import static com.android.systemui.pip.PipAnimationController.DURATION_DEFAULT_MS;
+
 import android.app.ActivityManager.RunningTaskInfo;
 import android.app.ActivityManager.StackInfo;
 import android.app.ActivityTaskManager;
@@ -50,7 +52,6 @@
 import com.android.systemui.UiOffloadThread;
 import com.android.systemui.broadcast.BroadcastDispatcher;
 import com.android.systemui.pip.BasePipManager;
-import com.android.systemui.pip.PipAnimationController;
 import com.android.systemui.pip.PipBoundsHandler;
 import com.android.systemui.pip.PipTaskOrganizer;
 import com.android.systemui.shared.system.ActivityManagerWrapper;
@@ -232,6 +233,7 @@
         if (mInitialized) {
             return;
         }
+
         mInitialized = true;
         mContext = context;
         mPipBoundsHandler = new PipBoundsHandler(context);
@@ -433,8 +435,7 @@
                 mCurrentPipBounds = mPipBounds;
                 break;
         }
-        mPipTaskOrganizer.animateResizePip(mCurrentPipBounds,
-                PipAnimationController.DURATION_DEFAULT_MS);
+        mPipTaskOrganizer.scheduleAnimateResizePip(mCurrentPipBounds, DURATION_DEFAULT_MS, null);
     }
 
     /**
diff --git a/packages/SystemUI/tests/src/com/android/systemui/pip/PipAnimationControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/pip/PipAnimationControllerTest.java
index 6c09a46..cd46110 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/pip/PipAnimationControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/pip/PipAnimationControllerTest.java
@@ -26,7 +26,6 @@
 import android.graphics.Rect;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
-import android.view.IWindowContainer;
 import android.view.SurfaceControl;
 
 import androidx.test.filters.SmallTest;
@@ -51,7 +50,7 @@
     private PipAnimationController mPipAnimationController;
 
     @Mock
-    private IWindowContainer mWindowContainer;
+    private SurfaceControl mLeash;
 
     @Mock
     private PipAnimationController.PipAnimationCallback mPipAnimationCallback;
@@ -65,8 +64,7 @@
     @Test
     public void getAnimator_withAlpha_returnFloatAnimator() {
         final PipAnimationController.PipTransitionAnimator animator = mPipAnimationController
-                .getAnimator(mWindowContainer, true /* scheduleFinishPip */,
-                        new Rect(), 0f, 1f);
+                .getAnimator(mLeash, true /* scheduleFinishPip */, new Rect(), 0f, 1f);
 
         assertEquals("Expect ANIM_TYPE_ALPHA animation",
                 animator.getAnimationType(), PipAnimationController.ANIM_TYPE_ALPHA);
@@ -75,8 +73,7 @@
     @Test
     public void getAnimator_withBounds_returnBoundsAnimator() {
         final PipAnimationController.PipTransitionAnimator animator = mPipAnimationController
-                .getAnimator(mWindowContainer, true /* scheduleFinishPip */,
-                        new Rect(), new Rect());
+                .getAnimator(mLeash, true /* scheduleFinishPip */, new Rect(), new Rect());
 
         assertEquals("Expect ANIM_TYPE_BOUNDS animation",
                 animator.getAnimationType(), PipAnimationController.ANIM_TYPE_BOUNDS);
@@ -88,14 +85,12 @@
         final Rect endValue1 = new Rect(100, 100, 200, 200);
         final Rect endValue2 = new Rect(200, 200, 300, 300);
         final PipAnimationController.PipTransitionAnimator oldAnimator = mPipAnimationController
-                .getAnimator(mWindowContainer, true /* scheduleFinishPip */,
-                        startValue, endValue1);
+                .getAnimator(mLeash, true /* scheduleFinishPip */, startValue, endValue1);
         oldAnimator.setSurfaceControlTransactionFactory(DummySurfaceControlTx::new);
         oldAnimator.start();
 
         final PipAnimationController.PipTransitionAnimator newAnimator = mPipAnimationController
-                .getAnimator(mWindowContainer, true /* scheduleFinishPip */,
-                        startValue, endValue2);
+                .getAnimator(mLeash, true /* scheduleFinishPip */, startValue, endValue2);
 
         assertEquals("getAnimator with same type returns same animator",
                 oldAnimator, newAnimator);
@@ -106,13 +101,11 @@
     @Test
     public void getAnimator_scheduleFinishPip() {
         PipAnimationController.PipTransitionAnimator animator = mPipAnimationController
-                .getAnimator(mWindowContainer, true /* scheduleFinishPip */,
-                        new Rect(), 0f, 1f);
+                .getAnimator(mLeash, true /* scheduleFinishPip */, new Rect(), 0f, 1f);
         assertTrue("scheduleFinishPip is true", animator.shouldScheduleFinishPip());
 
         animator = mPipAnimationController
-                .getAnimator(mWindowContainer, false /* scheduleFinishPip */,
-                        new Rect(), 0f, 1f);
+                .getAnimator(mLeash, false /* scheduleFinishPip */, new Rect(), 0f, 1f);
         assertFalse("scheduleFinishPip is false", animator.shouldScheduleFinishPip());
     }
 
@@ -122,8 +115,7 @@
         final Rect endValue1 = new Rect(100, 100, 200, 200);
         final Rect endValue2 = new Rect(200, 200, 300, 300);
         final PipAnimationController.PipTransitionAnimator animator = mPipAnimationController
-                .getAnimator(mWindowContainer, true /* scheduleFinishPip */,
-                        startValue, endValue1);
+                .getAnimator(mLeash, true /* scheduleFinishPip */, startValue, endValue1);
 
         animator.updateEndValue(endValue2);
 
@@ -135,24 +127,23 @@
         final Rect startValue = new Rect(0, 0, 100, 100);
         final Rect endValue = new Rect(100, 100, 200, 200);
         final PipAnimationController.PipTransitionAnimator animator = mPipAnimationController
-                .getAnimator(mWindowContainer, true /* scheduleFinishPip */,
-                        startValue, endValue);
+                .getAnimator(mLeash, true /* scheduleFinishPip */, startValue, endValue);
         animator.setSurfaceControlTransactionFactory(DummySurfaceControlTx::new);
 
         animator.setPipAnimationCallback(mPipAnimationCallback);
 
         // onAnimationStart triggers onPipAnimationStart
         animator.onAnimationStart(animator);
-        verify(mPipAnimationCallback).onPipAnimationStart(mWindowContainer, animator);
+        verify(mPipAnimationCallback).onPipAnimationStart(animator);
 
         // onAnimationCancel triggers onPipAnimationCancel
         animator.onAnimationCancel(animator);
-        verify(mPipAnimationCallback).onPipAnimationCancel(mWindowContainer, animator);
+        verify(mPipAnimationCallback).onPipAnimationCancel(animator);
 
         // onAnimationEnd triggers onPipAnimationEnd
         animator.onAnimationEnd(animator);
-        verify(mPipAnimationCallback).onPipAnimationEnd(eq(mWindowContainer),
-                any(SurfaceControl.Transaction.class), eq(animator));
+        verify(mPipAnimationCallback).onPipAnimationEnd(any(SurfaceControl.Transaction.class),
+                eq(animator));
     }
 
     /**