diff --git a/core/java/android/inputmethodservice/InputMethodService.java b/core/java/android/inputmethodservice/InputMethodService.java
index da7997d..fe14d457 100644
--- a/core/java/android/inputmethodservice/InputMethodService.java
+++ b/core/java/android/inputmethodservice/InputMethodService.java
@@ -974,6 +974,11 @@
                                 : InputMethodManager.RESULT_UNCHANGED_HIDDEN), null);
             }
             Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
+            if (android.view.inputmethod.Flags.refactorInsetsController()) {
+                // The hide request first finishes the animation and then proceeds to the server
+                // side, finally reaching here, marking this the end state.
+                ImeTracker.forLogging().onHidden(statsToken);
+            }
         }
 
         /**
@@ -3104,6 +3109,13 @@
 
         ImeTracker.forLogging().onProgress(statsToken, ImeTracker.PHASE_IME_SHOW_WINDOW);
 
+        if (android.view.inputmethod.Flags.refactorInsetsController()) {
+            // The ImeInsetsSourceProvider need the statsToken when dispatching the control
+            // (whenever the IME has drawn and its window is visible). Therefore, sending the
+            // statsToken here first.
+            notifyPreImeWindowVisibilityChanged(true /* visible */, statsToken);
+        }
+
         ImeTracing.getInstance().triggerServiceDump("InputMethodService#showWindow", mDumper,
                 null /* icProto */);
         Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "IMS.showWindow");
@@ -3127,7 +3139,9 @@
         if (DEBUG) Log.v(TAG, "showWindow: draw decorView!");
         mWindow.show();
         mDecorViewWasVisible = true;
-        applyVisibilityInInsetsConsumerIfNecessary(true /* setVisible */, statsToken);
+        if (!android.view.inputmethod.Flags.refactorInsetsController()) {
+            applyVisibilityInInsetsConsumerIfNecessary(true /* setVisible */, statsToken);
+        }
         cancelImeSurfaceRemoval();
         mInShowWindow = false;
         Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
@@ -3238,6 +3252,20 @@
                 ? mCurShowInputToken : mCurHideInputToken, setVisible, statsToken);
     }
 
+    /**
+     * Notifies the ImeInsetsSourceProvider before the IME visibility changes.
+     *
+     * @param visible {@code true} if it became visible, {@code false} otherwise.
+     * @param statsToken the token tracking the current IME request.
+     */
+    private void notifyPreImeWindowVisibilityChanged(boolean visible,
+            @NonNull ImeTracker.Token statsToken) {
+        final var viewRootImpl = getWindow().getWindow().getDecorView().getViewRootImpl();
+        if (viewRootImpl != null) {
+            viewRootImpl.notifyImeVisibilityChanged(visible, statsToken);
+        }
+    }
+
     private void finishViews(boolean finishingInput) {
         if (mInputViewStarted) {
             if (DEBUG) Log.v(TAG, "CALL: onFinishInputView");
@@ -3279,7 +3307,13 @@
         ImeTracing.getInstance().triggerServiceDump("InputMethodService#hideWindow", mDumper,
                 null /* icProto */);
         setImeWindowStatus(0, mBackDisposition);
-        applyVisibilityInInsetsConsumerIfNecessary(false /* setVisible */, statsToken);
+        if (android.view.inputmethod.Flags.refactorInsetsController()) {
+            // The ImeInsetsSourceProvider need the statsToken when dispatching the control. We
+            // send the token here, so that another request in the provider can be cancelled.
+            notifyPreImeWindowVisibilityChanged(false /* visible */, statsToken);
+        } else {
+            applyVisibilityInInsetsConsumerIfNecessary(false /* setVisible */, statsToken);
+        }
         mWindowVisible = false;
         finishViews(false /* finishingInput */);
         if (mDecorViewVisible) {
diff --git a/core/java/android/view/IWindowSession.aidl b/core/java/android/view/IWindowSession.aidl
index 14407ca..762a302 100644
--- a/core/java/android/view/IWindowSession.aidl
+++ b/core/java/android/view/IWindowSession.aidl
@@ -30,6 +30,7 @@
 import android.view.IWindowId;
 import android.view.MotionEvent;
 import android.view.WindowManager;
+import android.view.inputmethod.ImeTracker;
 import android.view.InsetsSourceControl;
 import android.view.InsetsState;
 import android.view.Surface;
@@ -276,7 +277,8 @@
     /**
      * Updates the requested visible types of insets.
      */
-    oneway void updateRequestedVisibleTypes(IWindow window, int requestedVisibleTypes);
+    oneway void updateRequestedVisibleTypes(IWindow window, int requestedVisibleTypes,
+            in @nullable ImeTracker.Token imeStatsToken);
 
     /**
      * Called when the system gesture exclusion has changed.
@@ -369,4 +371,14 @@
      * @return {@code true} if the focus changes. Otherwise, {@code false}.
      */
     boolean moveFocusToAdjacentWindow(IWindow fromWindow, int direction);
+
+    /**
+     * Notifies the statsToken and IME visibility to the ImeInsetsSourceProvider.
+     *
+     * @param window The window that is used to get the ImeInsetsSourceProvider.
+     * @param visible {@code true} to make it visible, {@code false} to hide it.
+     * @param statsToken the token tracking the current IME request.
+     */
+    oneway void notifyImeWindowVisibilityChangedFromClient(IWindow window, boolean visible,
+            in ImeTracker.Token statsToken);
 }
diff --git a/core/java/android/view/InsetsController.java b/core/java/android/view/InsetsController.java
index f166b89..d83f344 100644
--- a/core/java/android/view/InsetsController.java
+++ b/core/java/android/view/InsetsController.java
@@ -121,8 +121,10 @@
          * The visibilities should be reported back to WM.
          *
          * @param types Bitwise flags of types requested visible.
+         * @param statsToken the token tracking the current IME request or {@code null} otherwise.
          */
-        void updateRequestedVisibleTypes(@InsetsType int types);
+        void updateRequestedVisibleTypes(@InsetsType int types,
+                @Nullable ImeTracker.Token statsToken);
 
         /**
          * @return Whether the host has any callbacks it wants to synchronize the animations with.
@@ -974,6 +976,7 @@
         int consumedControlCount = 0;
         final @InsetsType int[] showTypes = new int[1];
         final @InsetsType int[] hideTypes = new int[1];
+        ImeTracker.Token statsToken = null;
 
         // Ensure to update all existing source consumers
         for (int i = mSourceConsumers.size() - 1; i >= 0; i--) {
@@ -988,6 +991,12 @@
             if (control != null) {
                 controllableTypes |= control.getType();
                 consumedControlCount++;
+
+                if (Flags.refactorInsetsController()) {
+                    if (control.getId() == ID_IME) {
+                        statsToken = control.getImeStatsToken();
+                    }
+                }
             }
 
             // control may be null, but we still need to update the control to null if it got
@@ -1021,34 +1030,31 @@
         if (Flags.refactorInsetsController()) {
             if (mPendingImeControlRequest != null && getImeSourceConsumer().getControl() != null
                     && getImeSourceConsumer().getControl().getLeash() != null) {
-                // TODO we need to pass the statsToken
-                handlePendingControlRequest(null);
+                handlePendingControlRequest(statsToken);
             } else {
                 if (showTypes[0] != 0) {
-                    applyAnimation(showTypes[0], true /* show */, false /* fromIme */,
-                            null /* statsToken */);
+                    applyAnimation(showTypes[0], true /* show */, false /* fromIme */, statsToken);
                 }
                 if (hideTypes[0] != 0) {
-                    applyAnimation(hideTypes[0], false /* show */, false /* fromIme */,
-                            null /* statsToken */);
+                    applyAnimation(hideTypes[0], false /* show */, false /* fromIme */, statsToken);
                 }
             }
         } else {
             if (showTypes[0] != 0) {
-                final var statsToken =
+                final var newStatsToken =
                         (showTypes[0] & ime()) == 0 ? null : ImeTracker.forLogging().onStart(
                                 ImeTracker.TYPE_SHOW, ImeTracker.ORIGIN_CLIENT,
                                 SoftInputShowHideReason.CONTROLS_CHANGED,
                                 mHost.isHandlingPointerEvent() /* fromUser */);
-                applyAnimation(showTypes[0], true /* show */, false /* fromIme */, statsToken);
+                applyAnimation(showTypes[0], true /* show */, false /* fromIme */, newStatsToken);
             }
             if (hideTypes[0] != 0) {
-                final var statsToken =
+                final var newStatsToken =
                         (hideTypes[0] & ime()) == 0 ? null : ImeTracker.forLogging().onStart(
                                 ImeTracker.TYPE_HIDE, ImeTracker.ORIGIN_CLIENT,
                                 SoftInputShowHideReason.CONTROLS_CHANGED,
                                 mHost.isHandlingPointerEvent() /* fromUser */);
-                applyAnimation(hideTypes[0], false /* show */, false /* fromIme */, statsToken);
+                applyAnimation(hideTypes[0], false /* show */, false /* fromIme */, newStatsToken);
             }
         }
 
@@ -1065,7 +1071,9 @@
         }
 
         // InsetsSourceConsumer#setControl might change the requested visibility.
-        reportRequestedVisibleTypes();
+        // TODO(b/353463205) check this: if the requestedVisibleTypes for the IME were already
+        //  sent, the request would fail. Therefore, don't send the statsToken here.
+        reportRequestedVisibleTypes(null /* statsToken */);
     }
 
     @VisibleForTesting(visibility = PACKAGE)
@@ -1176,6 +1184,7 @@
         }
         if (DEBUG) Log.d(TAG, "show typesReady: " + typesReady);
         if ((Flags.refactorInsetsController() || fromIme) && (typesReady & Type.ime()) != 0) {
+            // TODO(b/353463205) check if this is needed here
             ImeTracker.forLatency().onShown(statsToken, ActivityThread::currentApplication);
         }
         applyAnimation(typesReady, true /* show */, fromIme, statsToken);
@@ -1243,6 +1252,8 @@
                     // an animation again (mRequestedVisibleTypes are reported at the end of the IME
                     // hide animation but set at the beginning)
                     if ((mRequestedVisibleTypes & ime()) == 0) {
+                        ImeTracker.forLogging().onCancelled(statsToken,
+                                ImeTracker.PHASE_CLIENT_ALREADY_HIDDEN);
                         continue;
                     }
                 }
@@ -1346,7 +1357,7 @@
 
         // We are finishing setting the requested visible types. Report them to the server
         // and/or the app.
-        reportRequestedVisibleTypes();
+        reportRequestedVisibleTypes(statsToken);
     }
 
     private void controlAnimationUncheckedInner(@InsetsType int types,
@@ -1396,8 +1407,8 @@
             // Ime will not be contained in typesReady nor in controls, if we don't have a leash
             Pair<Integer, Integer> typesReadyPair = collectSourceControlsV2(types, controls);
             typesReady = typesReadyPair.first;
-            @InsetsType int typesWithoutLeash = typesReadyPair.second;
             if (animationType == ANIMATION_TYPE_USER) {
+                @InsetsType int typesWithoutLeash = typesReadyPair.second;
                 // When using an app-driven animation, the IME won't have a leash (because the
                 // window isn't created yet). If we have a control, but no leash, defers the
                 // request until the leash gets created.
@@ -1431,6 +1442,11 @@
                 }
                 // We need to wait until all types are ready
                 if (typesReady != types) {
+                    if (DEBUG) {
+                        Log.d(TAG, TextUtils.formatSimple(
+                                "not all types are ready yet, waiting. typesReady: %s, types: %s",
+                                typesReady, types));
+                    }
                     return;
                 }
             }
@@ -1728,9 +1744,13 @@
         } else {
             ImeTracker.forLogging().onProgress(statsToken,
                     ImeTracker.PHASE_CLIENT_ANIMATION_FINISHED_HIDE);
-            ImeTracker.forLogging().onHidden(statsToken);
+            // The requestedVisibleTypes are only send at the end of the hide animation.
+            // Therefore, the requested is not finished at this point.
+            if (!Flags.refactorInsetsController()) {
+                ImeTracker.forLogging().onHidden(statsToken);
+            }
         }
-        reportRequestedVisibleTypes();
+        reportRequestedVisibleTypes(shown ? null : runner.getStatsToken());
     }
 
     @Override
@@ -1787,7 +1807,7 @@
                             if (mHost != null) {
                                 // if the (hide) animation is cancelled, the
                                 // requestedVisibleTypes should be reported at this point.
-                                reportRequestedVisibleTypes();
+                                reportRequestedVisibleTypes(control.getStatsToken());
                                 mHost.getInputMethodManager().removeImeSurface(
                                         mHost.getWindowToken());
                             }
@@ -1923,8 +1943,10 @@
 
     /**
      * Called when finishing setting requested visible types or finishing setting controls.
+     *
+     * @param statsToken the token tracking the current IME request or {@code null} otherwise.
      */
-    private void reportRequestedVisibleTypes() {
+    private void reportRequestedVisibleTypes(@Nullable ImeTracker.Token statsToken) {
         final @InsetsType int typesToReport;
         if (Flags.refactorInsetsController()) {
             // If the IME is currently animating out, it is still visible, therefore we only
@@ -1941,8 +1963,23 @@
             if (WindowInsets.Type.hasCompatSystemBars(diff)) {
                 mCompatSysUiVisibilityStaled = true;
             }
+            if (Flags.refactorInsetsController()) {
+                ImeTracker.forLogging().onProgress(statsToken,
+                        ImeTracker.PHASE_CLIENT_REPORT_REQUESTED_VISIBLE_TYPES);
+            }
             mReportedRequestedVisibleTypes = mRequestedVisibleTypes;
-            mHost.updateRequestedVisibleTypes(mReportedRequestedVisibleTypes);
+            mHost.updateRequestedVisibleTypes(mReportedRequestedVisibleTypes, statsToken);
+        } else if (Flags.refactorInsetsController()) {
+            if ((typesToReport & ime()) != 0 && mImeSourceConsumer != null) {
+                InsetsSourceControl control = mImeSourceConsumer.getControl();
+                if (control != null && control.getLeash() == null) {
+                    // If the IME was requested twice, and we didn't receive the controls
+                    // yet, this request will not continue. It should be cancelled here, as
+                    // it would time out otherwise.
+                    ImeTracker.forLogging().onCancelled(statsToken,
+                            ImeTracker.PHASE_CLIENT_REPORT_REQUESTED_VISIBLE_TYPES);
+                }
+            }
         }
         updateCompatSysUiVisibility();
     }
diff --git a/core/java/android/view/InsetsSourceControl.java b/core/java/android/view/InsetsSourceControl.java
index 2efa647..7877352 100644
--- a/core/java/android/view/InsetsSourceControl.java
+++ b/core/java/android/view/InsetsSourceControl.java
@@ -31,6 +31,7 @@
 import android.os.Parcelable;
 import android.util.proto.ProtoOutputStream;
 import android.view.WindowInsets.Type.InsetsType;
+import android.view.inputmethod.ImeTracker;
 
 import java.io.PrintWriter;
 import java.util.Arrays;
@@ -56,6 +57,9 @@
     private boolean mSkipAnimationOnce;
     private int mParcelableFlags;
 
+    /** The token tracking the current IME request */
+    private @Nullable ImeTracker.Token mImeStatsToken;
+
     public InsetsSourceControl(int id, @InsetsType int type, @Nullable SurfaceControl leash,
             boolean initiallyVisible, Point surfacePosition, Insets insetsHint) {
         mId = id;
@@ -78,6 +82,7 @@
         mSurfacePosition = new Point(other.mSurfacePosition);
         mInsetsHint = other.mInsetsHint;
         mSkipAnimationOnce = other.getAndClearSkipAnimationOnce();
+        mImeStatsToken = other.getImeStatsToken();
     }
 
     public InsetsSourceControl(Parcel in) {
@@ -88,6 +93,7 @@
         mSurfacePosition = in.readTypedObject(Point.CREATOR);
         mInsetsHint = in.readTypedObject(Insets.CREATOR);
         mSkipAnimationOnce = in.readBoolean();
+        mImeStatsToken = in.readTypedObject(ImeTracker.Token.CREATOR);
     }
 
     public int getId() {
@@ -153,6 +159,15 @@
         return result;
     }
 
+    @Nullable
+    public ImeTracker.Token getImeStatsToken() {
+        return mImeStatsToken;
+    }
+
+    public void setImeStatsToken(@Nullable ImeTracker.Token imeStatsToken) {
+        mImeStatsToken = imeStatsToken;
+    }
+
     public void setParcelableFlags(int parcelableFlags) {
         mParcelableFlags = parcelableFlags;
     }
@@ -171,6 +186,7 @@
         dest.writeTypedObject(mSurfacePosition, mParcelableFlags);
         dest.writeTypedObject(mInsetsHint, mParcelableFlags);
         dest.writeBoolean(mSkipAnimationOnce);
+        dest.writeTypedObject(mImeStatsToken, mParcelableFlags);
     }
 
     public void release(Consumer<SurfaceControl> surfaceReleaseConsumer) {
@@ -196,13 +212,14 @@
                 && mInitiallyVisible == that.mInitiallyVisible
                 && mSurfacePosition.equals(that.mSurfacePosition)
                 && mInsetsHint.equals(that.mInsetsHint)
-                && mSkipAnimationOnce == that.mSkipAnimationOnce;
+                && mSkipAnimationOnce == that.mSkipAnimationOnce
+                && Objects.equals(mImeStatsToken, that.mImeStatsToken);
     }
 
     @Override
     public int hashCode() {
         return Objects.hash(mId, mType, mLeash, mInitiallyVisible, mSurfacePosition, mInsetsHint,
-                mSkipAnimationOnce);
+                mSkipAnimationOnce, mImeStatsToken);
     }
 
     @Override
@@ -225,6 +242,7 @@
         pw.print(" mSurfacePosition="); pw.print(mSurfacePosition);
         pw.print(" mInsetsHint="); pw.print(mInsetsHint);
         pw.print(" mSkipAnimationOnce="); pw.print(mSkipAnimationOnce);
+        pw.print(" mImeStatsToken="); pw.print(mImeStatsToken);
         pw.println();
     }
 
diff --git a/core/java/android/view/ViewRootImpl.java b/core/java/android/view/ViewRootImpl.java
index 2f204f9..37d5220 100644
--- a/core/java/android/view/ViewRootImpl.java
+++ b/core/java/android/view/ViewRootImpl.java
@@ -10057,6 +10057,24 @@
     }
 
     /**
+     * Dispatches the statsToken and IME visibility to the ImeInsetsSourceProvider.
+     *
+     * @param visible {@code true} if it became visible, {@code false} otherwise.
+     * @param statsToken the token tracking the current IME request.
+     *
+     * @hide
+     */
+    public void notifyImeVisibilityChanged(boolean visible, @NonNull ImeTracker.Token statsToken) {
+        ImeTracker.forLogging().onProgress(statsToken,
+                ImeTracker.PHASE_CLIENT_NOTIFY_IME_VISIBILITY_CHANGED);
+        try {
+            mWindowSession.notifyImeWindowVisibilityChangedFromClient(mWindow, visible, statsToken);
+        } catch (RemoteException e) {
+            e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
      * Represents a pending input event that is waiting in a queue.
      *
      * Input events are processed in serial order by the timestamp specified by
diff --git a/core/java/android/view/ViewRootInsetsControllerHost.java b/core/java/android/view/ViewRootInsetsControllerHost.java
index b66c59a..889acca4 100644
--- a/core/java/android/view/ViewRootInsetsControllerHost.java
+++ b/core/java/android/view/ViewRootInsetsControllerHost.java
@@ -19,12 +19,14 @@
 import static android.view.InsetsController.DEBUG;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.content.Context;
 import android.content.res.CompatibilityInfo;
 import android.os.Handler;
 import android.os.IBinder;
 import android.os.RemoteException;
 import android.util.Log;
+import android.view.inputmethod.ImeTracker;
 import android.view.inputmethod.InputMethodManager;
 
 import java.util.List;
@@ -151,10 +153,17 @@
     }
 
     @Override
-    public void updateRequestedVisibleTypes(@WindowInsets.Type.InsetsType int types) {
+    public void updateRequestedVisibleTypes(@WindowInsets.Type.InsetsType int types,
+            @Nullable ImeTracker.Token statsToken) {
         try {
             if (mViewRoot.mAdded) {
-                mViewRoot.mWindowSession.updateRequestedVisibleTypes(mViewRoot.mWindow, types);
+                ImeTracker.forLogging().onProgress(statsToken,
+                        ImeTracker.PHASE_CLIENT_UPDATE_REQUESTED_VISIBLE_TYPES);
+                mViewRoot.mWindowSession.updateRequestedVisibleTypes(mViewRoot.mWindow, types,
+                        statsToken);
+            } else {
+                ImeTracker.forLogging().onFailed(statsToken,
+                        ImeTracker.PHASE_CLIENT_UPDATE_REQUESTED_VISIBLE_TYPES);
             }
         } catch (RemoteException e) {
             Log.e(TAG, "Failed to call insetsModified", e);
diff --git a/core/java/android/view/WindowlessWindowManager.java b/core/java/android/view/WindowlessWindowManager.java
index 7871858..0d027f1 100644
--- a/core/java/android/view/WindowlessWindowManager.java
+++ b/core/java/android/view/WindowlessWindowManager.java
@@ -30,6 +30,7 @@
 import android.util.MergedConfiguration;
 import android.view.View.FocusDirection;
 import android.view.WindowInsets.Type.InsetsType;
+import android.view.inputmethod.ImeTracker;
 import android.window.ClientWindowFrames;
 import android.window.InputTransferToken;
 import android.window.OnBackInvokedCallbackInfo;
@@ -596,7 +597,7 @@
 
     @Override
     public void updateRequestedVisibleTypes(IWindow window,
-            @InsetsType int requestedVisibleTypes)  {
+            @InsetsType int requestedVisibleTypes, @Nullable ImeTracker.Token imeStatsToken)  {
     }
 
     @Override
@@ -677,6 +678,11 @@
         return false;
     }
 
+    @Override
+    public void notifyImeWindowVisibilityChangedFromClient(IWindow window, boolean visible,
+            @NonNull ImeTracker.Token statsToken) {
+    }
+
     void setParentInterface(@Nullable ISurfaceControlViewHostParent parentInterface) {
         IBinder oldInterface = mParentInterface == null ? null : mParentInterface.asBinder();
         IBinder newInterface = parentInterface == null ? null : parentInterface.asBinder();
diff --git a/core/java/android/view/inputmethod/ImeTracker.java b/core/java/android/view/inputmethod/ImeTracker.java
index edc9921..b9751c8 100644
--- a/core/java/android/view/inputmethod/ImeTracker.java
+++ b/core/java/android/view/inputmethod/ImeTracker.java
@@ -202,6 +202,24 @@
             PHASE_IME_HIDE_WINDOW,
             PHASE_IME_PRIVILEGED_OPERATIONS,
             PHASE_SERVER_CURRENT_ACTIVE_IME,
+            PHASE_CLIENT_REPORT_REQUESTED_VISIBLE_TYPES,
+            PHASE_WM_SET_REMOTE_TARGET_IME_VISIBILITY,
+            PHASE_WM_POST_LAYOUT_NOTIFY_CONTROLS_CHANGED,
+            PHASE_CLIENT_HANDLE_DISPATCH_IME_VISIBILITY_CHANGED,
+            PHASE_CLIENT_NOTIFY_IME_VISIBILITY_CHANGED,
+            PHASE_CLIENT_UPDATE_REQUESTED_VISIBLE_TYPES,
+            PHASE_WM_REMOTE_INSETS_CONTROL_TARGET_SET_REQUESTED_VISIBILITY,
+            PHASE_WM_GET_CONTROL_WITH_LEASH,
+            PHASE_WM_UPDATE_REQUESTED_VISIBLE_TYPES,
+            PHASE_SERVER_SET_VISIBILITY_ON_FOCUSED_WINDOW,
+            PHASE_CLIENT_HANDLE_SET_IME_VISIBILITY,
+            PHASE_CLIENT_SET_IME_VISIBILITY,
+            PHASE_WM_DISPATCH_IME_REQUESTED_CHANGED,
+            PHASE_CLIENT_NO_ONGOING_USER_ANIMATION,
+            PHASE_WM_NOTIFY_IME_VISIBILITY_CHANGED_FROM_CLIENT,
+            PHASE_WM_POSTING_CHANGED_IME_VISIBILITY,
+            PHASE_WM_INVOKING_IME_REQUESTED_LISTENER,
+            PHASE_CLIENT_ALREADY_HIDDEN,
     })
     @Retention(RetentionPolicy.SOURCE)
     @interface Phase {}
@@ -351,6 +369,62 @@
     /** Checked that the calling IME is the currently active IME. */
     int PHASE_SERVER_CURRENT_ACTIVE_IME = ImeProtoEnums.PHASE_SERVER_CURRENT_ACTIVE_IME;
 
+    /** Reporting the new requested visible types. */
+    int PHASE_CLIENT_REPORT_REQUESTED_VISIBLE_TYPES =
+            ImeProtoEnums.PHASE_CLIENT_REPORT_REQUESTED_VISIBLE_TYPES;
+    /** Setting the IME visibility for the RemoteInsetsControlTarget. */
+    int PHASE_WM_SET_REMOTE_TARGET_IME_VISIBILITY =
+            ImeProtoEnums.PHASE_WM_SET_REMOTE_TARGET_IME_VISIBILITY;
+    /** IME has no insets pending and is server visible. Notify about changed controls. */
+    int PHASE_WM_POST_LAYOUT_NOTIFY_CONTROLS_CHANGED =
+            ImeProtoEnums.PHASE_WM_POST_LAYOUT_NOTIFY_CONTROLS_CHANGED;
+    /** Handling the dispatch of the IME visibility change. */
+    int PHASE_CLIENT_HANDLE_DISPATCH_IME_VISIBILITY_CHANGED =
+            ImeProtoEnums.PHASE_CLIENT_HANDLE_DISPATCH_IME_VISIBILITY_CHANGED;
+    /** Dispatching the IME visibility change. */
+    int PHASE_CLIENT_NOTIFY_IME_VISIBILITY_CHANGED =
+            ImeProtoEnums.PHASE_CLIENT_NOTIFY_IME_VISIBILITY_CHANGED;
+    /** Updating the requested visible types. */
+    int PHASE_CLIENT_UPDATE_REQUESTED_VISIBLE_TYPES =
+            ImeProtoEnums.PHASE_CLIENT_UPDATE_REQUESTED_VISIBLE_TYPES;
+    /** Reached the remote insets control target's setImeInputTargetRequestedVisibility method. */
+    int PHASE_WM_REMOTE_INSETS_CONTROL_TARGET_SET_REQUESTED_VISIBILITY =
+            ImeProtoEnums.PHASE_WM_REMOTE_INSETS_CONTROL_TARGET_SET_REQUESTED_VISIBILITY;
+    /** Received a new insets source control with a leash. */
+    int PHASE_WM_GET_CONTROL_WITH_LEASH =
+            ImeProtoEnums.PHASE_WM_GET_CONTROL_WITH_LEASH;
+    /**
+     * Updating the requested visible types in the WindowState and sending them to state
+     * controller.
+     */
+    int PHASE_WM_UPDATE_REQUESTED_VISIBLE_TYPES =
+            ImeProtoEnums.PHASE_WM_UPDATE_REQUESTED_VISIBLE_TYPES;
+    /** Setting the requested IME visibility of a window. */
+    int PHASE_SERVER_SET_VISIBILITY_ON_FOCUSED_WINDOW =
+            ImeProtoEnums.PHASE_SERVER_SET_VISIBILITY_ON_FOCUSED_WINDOW;
+    /** Reached the redirect of InputMethodManager to InsetsController show/hide. */
+    int PHASE_CLIENT_HANDLE_SET_IME_VISIBILITY =
+            ImeProtoEnums.PHASE_CLIENT_HANDLE_SET_IME_VISIBILITY;
+    /** Reached the InputMethodManager Handler call to send the visibility. */
+    int PHASE_CLIENT_SET_IME_VISIBILITY = ImeProtoEnums.PHASE_CLIENT_SET_IME_VISIBILITY;
+    /** Calling into the listener to show/hide the IME from the ImeInsetsSourceProvider. */
+    int PHASE_WM_DISPATCH_IME_REQUESTED_CHANGED =
+            ImeProtoEnums.PHASE_WM_DISPATCH_IME_REQUESTED_CHANGED;
+    /** An ongoing user animation will not be interrupted by a IMM#showSoftInput. */
+    int PHASE_CLIENT_NO_ONGOING_USER_ANIMATION =
+            ImeProtoEnums.PHASE_CLIENT_NO_ONGOING_USER_ANIMATION;
+    /** Dispatching the token to the ImeInsetsSourceProvider. */
+    int PHASE_WM_NOTIFY_IME_VISIBILITY_CHANGED_FROM_CLIENT =
+            ImeProtoEnums.PHASE_WM_NOTIFY_IME_VISIBILITY_CHANGED_FROM_CLIENT;
+    /** Now posting the IME visibility to the WMS handler. */
+    int PHASE_WM_POSTING_CHANGED_IME_VISIBILITY =
+            ImeProtoEnums.PHASE_WM_POSTING_CHANGED_IME_VISIBILITY;
+    /** Inside the WMS handler calling into the listener that calls into IMMS show/hide. */
+    int PHASE_WM_INVOKING_IME_REQUESTED_LISTENER =
+            ImeProtoEnums.PHASE_WM_INVOKING_IME_REQUESTED_LISTENER;
+    /** IME is requested to be hidden, but already hidden. Don't hide to avoid another animation. */
+    int PHASE_CLIENT_ALREADY_HIDDEN = ImeProtoEnums.PHASE_CLIENT_ALREADY_HIDDEN;
+
     /**
      * Called when an IME request is started.
      *
diff --git a/core/java/android/view/inputmethod/InputMethodManager.java b/core/java/android/view/inputmethod/InputMethodManager.java
index fed8eea..d40b72c 100644
--- a/core/java/android/view/inputmethod/InputMethodManager.java
+++ b/core/java/android/view/inputmethod/InputMethodManager.java
@@ -1350,11 +1350,15 @@
                     return;
                 }
                 case MSG_SET_VISIBILITY:
-                    final boolean visible = msg.arg1 != 0;
+                    final SomeArgs args = (SomeArgs) msg.obj;
+                    final boolean visible = (boolean) args.arg1;
+                    final ImeTracker.Token statsToken = (ImeTracker.Token) args.arg2;
                     synchronized (mH) {
                         if (mCurRootView != null) {
                             final var insetsController = mCurRootView.getInsetsController();
                             if (insetsController != null) {
+                                ImeTracker.forLogging().onProgress(statsToken,
+                                        ImeTracker.PHASE_CLIENT_HANDLE_SET_IME_VISIBILITY);
                                 if (visible) {
                                     insetsController.show(WindowInsets.Type.ime(),
                                             false /* fromIme */, null /* statsToken */);
@@ -1363,6 +1367,9 @@
                                             false /* fromIme */, null /* statsToken */);
                                 }
                             }
+                        } else {
+                            ImeTracker.forLogging().onFailed(statsToken,
+                                    ImeTracker.PHASE_CLIENT_HANDLE_SET_IME_VISIBILITY);
                         }
                     }
                     break;
@@ -1463,8 +1470,13 @@
         }
 
         @Override
-        public void setImeVisibility(boolean visible) {
-            mH.obtainMessage(MSG_SET_VISIBILITY, visible ? 1 : 0, 0).sendToTarget();
+        public void setImeVisibility(boolean visible, @Nullable ImeTracker.Token statsToken) {
+            final SomeArgs args = SomeArgs.obtain();
+            args.arg1 = visible;
+            args.arg2 = statsToken;
+            ImeTracker.forLogging().onProgress(statsToken,
+                    ImeTracker.PHASE_CLIENT_SET_IME_VISIBILITY);
+            mH.obtainMessage(MSG_SET_VISIBILITY, args).sendToTarget();
         }
 
         @Override
@@ -2344,11 +2356,15 @@
                 if (viewRootImpl != null
                         && (viewRootImpl.getInsetsController().computeUserAnimatingTypes()
                                 & WindowInsets.Type.ime()) == 0) {
+                    ImeTracker.forLogging().onProgress(statsToken,
+                            ImeTracker.PHASE_CLIENT_NO_ONGOING_USER_ANIMATION);
                     // TODO(b/322992891) handle case of SHOW_IMPLICIT
                     viewRootImpl.getInsetsController().show(WindowInsets.Type.ime(),
                             false /* fromIme */, statsToken);
                     return true;
                 }
+                ImeTracker.forLogging().onCancelled(statsToken,
+                        ImeTracker.PHASE_CLIENT_NO_ONGOING_USER_ANIMATION);
                 return false;
             } else {
                 // Makes sure to call ImeInsetsSourceConsumer#onShowRequested on the UI thread.
diff --git a/core/java/com/android/internal/inputmethod/IInputMethodClient.aidl b/core/java/com/android/internal/inputmethod/IInputMethodClient.aidl
index 6a7fa99..6eb9d2e 100644
--- a/core/java/com/android/internal/inputmethod/IInputMethodClient.aidl
+++ b/core/java/com/android/internal/inputmethod/IInputMethodClient.aidl
@@ -16,6 +16,7 @@
 
 package com.android.internal.inputmethod;
 
+import android.view.inputmethod.ImeTracker;
 import com.android.internal.inputmethod.InputBindResult;
 
 /**
@@ -30,7 +31,7 @@
     void onUnbindAccessibilityService(int sequence, int id);
     void setActive(boolean active, boolean fullscreen);
     void setInteractive(boolean active, boolean fullscreen);
-    void setImeVisibility(boolean visible);
+    void setImeVisibility(boolean visible, in @nullable ImeTracker.Token statsToken);
     void scheduleStartInputIfNecessary(boolean fullscreen);
     void reportFullscreenMode(boolean fullscreen);
     void setImeTraceEnabled(boolean enabled);
diff --git a/core/tests/coretests/src/android/view/InsetsControllerTest.java b/core/tests/coretests/src/android/view/InsetsControllerTest.java
index ae7f465..7bc0d2f 100644
--- a/core/tests/coretests/src/android/view/InsetsControllerTest.java
+++ b/core/tests/coretests/src/android/view/InsetsControllerTest.java
@@ -56,6 +56,7 @@
 import static java.util.concurrent.TimeUnit.SECONDS;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.content.Context;
 import android.graphics.Insets;
 import android.graphics.Point;
@@ -1095,9 +1096,10 @@
         }
 
         @Override
-        public void updateRequestedVisibleTypes(@InsetsType int requestedVisibleTypes) {
+        public void updateRequestedVisibleTypes(@InsetsType int requestedVisibleTypes,
+                @Nullable ImeTracker.Token statsToken) {
             mRequestedVisibleTypes = requestedVisibleTypes;
-            super.updateRequestedVisibleTypes(requestedVisibleTypes);
+            super.updateRequestedVisibleTypes(requestedVisibleTypes, statsToken);
         }
 
         public boolean isRequestedVisible(@InsetsType int types) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java
index f4ac5f2..4fbb574 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java
@@ -344,8 +344,7 @@
 
             if (android.view.inputmethod.Flags.refactorInsetsController()) {
                 if (pendingImeStartAnimation) {
-                    startAnimation(true, true /* forceRestart */,
-                            null /* statsToken */);
+                    startAnimation(true, true /* forceRestart */);
                 }
             }
         }
@@ -398,8 +397,7 @@
                 // already (e.g., when focussing an editText in activity B, while and editText in
                 // activity A is focussed), we will not get a call of #insetsControlChanged, and
                 // therefore have to start the show animation from here
-                startAnimation(mImeRequestedVisible /* show */, false /* forceRestart */,
-                        null /* TODO statsToken */);
+                startAnimation(mImeRequestedVisible /* show */, false /* forceRestart */);
             }
         }
 
@@ -436,16 +434,31 @@
                     .navBarFrameHeight();
         }
 
+        private void startAnimation(final boolean show, final boolean forceRestart) {
+            final var imeSource = mInsetsState.peekSource(InsetsSource.ID_IME);
+            if (imeSource == null || mImeSourceControl == null) {
+                return;
+            }
+            final var statsToken = mImeSourceControl.getImeStatsToken();
+
+            startAnimation(show, forceRestart, statsToken);
+        }
+
         private void startAnimation(final boolean show, final boolean forceRestart,
                 @SoftInputShowHideReason int reason) {
             final var imeSource = mInsetsState.peekSource(InsetsSource.ID_IME);
             if (imeSource == null || mImeSourceControl == null) {
                 return;
             }
-            final var statsToken = ImeTracker.forLogging().onStart(
-                    show ? ImeTracker.TYPE_SHOW : ImeTracker.TYPE_HIDE, ImeTracker.ORIGIN_WM_SHELL,
-                    reason, false /* fromUser */);
-
+            final ImeTracker.Token statsToken;
+            if (android.view.inputmethod.Flags.refactorInsetsController()
+                    && mImeSourceControl.getImeStatsToken() != null) {
+                statsToken = mImeSourceControl.getImeStatsToken();
+            } else {
+                statsToken = ImeTracker.forLogging().onStart(
+                        show ? ImeTracker.TYPE_SHOW : ImeTracker.TYPE_HIDE,
+                        ImeTracker.ORIGIN_WM_SHELL, reason, false /* fromUser */);
+            }
             startAnimation(show, forceRestart, statsToken);
         }
 
diff --git a/services/core/java/com/android/server/inputmethod/DefaultImeVisibilityApplier.java b/services/core/java/com/android/server/inputmethod/DefaultImeVisibilityApplier.java
index 7c93c8b..5c939bc 100644
--- a/services/core/java/com/android/server/inputmethod/DefaultImeVisibilityApplier.java
+++ b/services/core/java/com/android/server/inputmethod/DefaultImeVisibilityApplier.java
@@ -202,7 +202,7 @@
                 break;
             case STATE_HIDE_IME_EXPLICIT:
                 if (Flags.refactorInsetsController()) {
-                    setImeVisibilityOnFocusedWindowClient(false, userId);
+                    setImeVisibilityOnFocusedWindowClient(false, userId, statsToken);
                 } else {
                     mService.hideCurrentInputLocked(windowToken, statsToken,
                             0 /* flags */, null /* resultReceiver */, reason, userId);
@@ -210,7 +210,7 @@
                 break;
             case STATE_HIDE_IME_NOT_ALWAYS:
                 if (Flags.refactorInsetsController()) {
-                    setImeVisibilityOnFocusedWindowClient(false, userId);
+                    setImeVisibilityOnFocusedWindowClient(false, userId, statsToken);
                 } else {
                     mService.hideCurrentInputLocked(windowToken, statsToken,
                             InputMethodManager.HIDE_NOT_ALWAYS, null /* resultReceiver */, reason,
@@ -221,7 +221,7 @@
                 if (Flags.refactorInsetsController()) {
                     // This can be triggered by IMMS#startInputOrWindowGainedFocus. We need to
                     // set the requestedVisibleTypes in InsetsController first, before applying it.
-                    setImeVisibilityOnFocusedWindowClient(true, userId);
+                    setImeVisibilityOnFocusedWindowClient(true, userId, statsToken);
                 } else {
                     mService.showCurrentInputLocked(windowToken, statsToken,
                             InputMethodManager.SHOW_IMPLICIT, MotionEvent.TOOL_TYPE_UNKNOWN,
@@ -278,14 +278,17 @@
     }
 
     @GuardedBy("ImfLock.class")
-    private void setImeVisibilityOnFocusedWindowClient(boolean visibility, @UserIdInt int userId) {
+    private void setImeVisibilityOnFocusedWindowClient(boolean visibility, @UserIdInt int userId,
+            @NonNull ImeTracker.Token statsToken) {
         final var userData = mService.getUserData(userId);
         if (userData.mImeBindingState != null
                 && userData.mImeBindingState.mFocusedWindowClient != null
                 && userData.mImeBindingState.mFocusedWindowClient.mClient != null) {
-            userData.mImeBindingState.mFocusedWindowClient.mClient.setImeVisibility(visibility);
+            userData.mImeBindingState.mFocusedWindowClient.mClient.setImeVisibility(visibility,
+                    statsToken);
         } else {
-            // TODO(b/329229469): ImeTracker?
+            ImeTracker.forLogging().onFailed(statsToken,
+                    ImeTracker.PHASE_SERVER_SET_VISIBILITY_ON_FOCUSED_WINDOW);
         }
     }
 }
diff --git a/services/core/java/com/android/server/inputmethod/IInputMethodClientInvoker.java b/services/core/java/com/android/server/inputmethod/IInputMethodClientInvoker.java
index eada288..650ea60 100644
--- a/services/core/java/com/android/server/inputmethod/IInputMethodClientInvoker.java
+++ b/services/core/java/com/android/server/inputmethod/IInputMethodClientInvoker.java
@@ -25,6 +25,7 @@
 import android.os.IBinder;
 import android.os.RemoteException;
 import android.util.Slog;
+import android.view.inputmethod.ImeTracker;
 
 import com.android.internal.inputmethod.IInputMethodClient;
 import com.android.internal.inputmethod.InputBindResult;
@@ -250,18 +251,18 @@
     }
 
     @AnyThread
-    void setImeVisibility(boolean visible) {
+    void setImeVisibility(boolean visible, @Nullable ImeTracker.Token statsToken) {
         if (mIsProxy) {
-            setImeVisibilityInternal(visible);
+            setImeVisibilityInternal(visible, statsToken);
         } else {
-            mHandler.post(() -> setImeVisibilityInternal(visible));
+            mHandler.post(() -> setImeVisibilityInternal(visible, statsToken));
         }
     }
 
     @AnyThread
-    private void setImeVisibilityInternal(boolean visible) {
+    private void setImeVisibilityInternal(boolean visible, @Nullable ImeTracker.Token statsToken) {
         try {
-            mTarget.setImeVisibility(visible);
+            mTarget.setImeVisibility(visible, statsToken);
         } catch (RemoteException e) {
             logRemoteException(e);
         }
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
index a9723cc..f93e458 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
@@ -1429,12 +1429,12 @@
 
     void registerImeRequestedChangedListener() {
         mWindowManagerInternal.setOnImeRequestedChangedListener(
-                (windowToken, imeVisible) -> {
+                (windowToken, imeVisible, statsToken) -> {
                     if (Flags.refactorInsetsController()) {
                         if (imeVisible) {
-                            showCurrentInputInternal(windowToken);
+                            showCurrentInputInternal(windowToken, statsToken);
                         } else {
-                            hideCurrentInputInternal(windowToken);
+                            hideCurrentInputInternal(windowToken, statsToken);
                         }
                     }
                 });
@@ -1887,7 +1887,12 @@
         if (Flags.refactorInsetsController()) {
             if (isShowRequestedForCurrentWindow(userId) && userData.mImeBindingState != null
                     && userData.mImeBindingState.mFocusedWindow != null) {
-                showCurrentInputInternal(userData.mImeBindingState.mFocusedWindow);
+                // Re-use current statsToken, if it exists.
+                final var statsToken = userData.mCurStatsToken != null ? userData.mCurStatsToken
+                        : createStatsTokenForFocusedClient(true /* show */,
+                                SoftInputShowHideReason.ATTACH_NEW_INPUT, userId);
+                userData.mCurStatsToken = null;
+                showCurrentInputInternal(userData.mImeBindingState.mFocusedWindow, statsToken);
             }
         } else {
             if (isShowRequestedForCurrentWindow(userId)) {
@@ -3099,7 +3104,7 @@
                             && userData.mImeBindingState.mFocusedWindowClient != null
                             && userData.mImeBindingState.mFocusedWindowClient.mClient != null) {
                         userData.mImeBindingState.mFocusedWindowClient.mClient
-                                .setImeVisibility(true);
+                                .setImeVisibility(true, statsToken);
                         if (resultReceiver != null) {
                             resultReceiver.send(
                                     wasVisible ? InputMethodManager.RESULT_UNCHANGED_SHOWN
@@ -3119,7 +3124,8 @@
         }
     }
 
-    boolean showCurrentInputInternal(IBinder windowToken) {
+    // TODO(b/353463205) check callers to see if we can make statsToken @NonNull
+    boolean showCurrentInputInternal(IBinder windowToken, @Nullable ImeTracker.Token statsToken) {
         Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "IMMS.showCurrentInputInternal");
         ImeTracing.getInstance().triggerManagerServiceDump(
                 "InputMethodManagerService#showSoftInput", mDumper);
@@ -3129,7 +3135,7 @@
             final long ident = Binder.clearCallingIdentity();
             try {
                 if (DEBUG) Slog.v(TAG, "Client requesting input be shown");
-                return showCurrentInputLocked(windowToken, null /* statsToken */, 0 /* flags */,
+                return showCurrentInputLocked(windowToken, statsToken, 0 /* flags */,
                         0 /* lastClickTooType */, null /* resultReceiver */,
                         SoftInputShowHideReason.SHOW_SOFT_INPUT, userId);
             } finally {
@@ -3139,7 +3145,8 @@
         }
     }
 
-    boolean hideCurrentInputInternal(IBinder windowToken) {
+    // TODO(b/353463205) check callers to see if we can make statsToken @NonNull
+    boolean hideCurrentInputInternal(IBinder windowToken, @Nullable ImeTracker.Token statsToken) {
         Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "IMMS.hideCurrentInputInternal");
         ImeTracing.getInstance().triggerManagerServiceDump(
                 "InputMethodManagerService#hideSoftInput", mDumper);
@@ -3149,7 +3156,7 @@
             final long ident = Binder.clearCallingIdentity();
             try {
                 if (DEBUG) Slog.v(TAG, "Client requesting input be hidden");
-                return hideCurrentInputLocked(windowToken, null /* statsToken */, 0 /* flags */,
+                return hideCurrentInputLocked(windowToken, statsToken, 0 /* flags */,
                         null /* resultReceiver */, SoftInputShowHideReason.HIDE_SOFT_INPUT,
                         userId);
             } finally {
@@ -3536,7 +3543,7 @@
                         boolean wasVisible = isInputShownLocked();
                         // TODO add windowToken to interface
                         userData.mImeBindingState.mFocusedWindowClient.mClient
-                                .setImeVisibility(false);
+                                .setImeVisibility(false, statsToken);
                         if (resultReceiver != null) {
                             resultReceiver.send(wasVisible ? InputMethodManager.RESULT_HIDDEN
                                     : InputMethodManager.RESULT_UNCHANGED_HIDDEN, null);
@@ -4855,13 +4862,13 @@
                 final long ident = Binder.clearCallingIdentity();
                 try {
                     if (Flags.refactorInsetsController()) {
-                        userData.mCurClient.mClient.setImeVisibility(false);
+                        userData.mCurClient.mClient.setImeVisibility(false, statsToken);
                         // TODO we will loose the flags here
                         if (userData.mImeBindingState != null
                                 && userData.mImeBindingState.mFocusedWindowClient != null
                                 && userData.mImeBindingState.mFocusedWindowClient.mClient != null) {
                             userData.mImeBindingState.mFocusedWindowClient.mClient
-                                    .setImeVisibility(false);
+                                    .setImeVisibility(false, statsToken);
                         }
                     } else {
                         hideCurrentInputLocked(mLastImeTargetWindow, statsToken, flags,
@@ -4894,13 +4901,13 @@
                 final long ident = Binder.clearCallingIdentity();
                 try {
                     if (Flags.refactorInsetsController()) {
-                        userData.mCurClient.mClient.setImeVisibility(false);
+                        userData.mCurClient.mClient.setImeVisibility(false, statsToken);
                         // TODO we will loose the flags here
                         if (userData.mImeBindingState != null
                                 && userData.mImeBindingState.mFocusedWindowClient != null
                                 && userData.mImeBindingState.mFocusedWindowClient.mClient != null) {
                             userData.mImeBindingState.mFocusedWindowClient.mClient
-                                    .setImeVisibility(true);
+                                    .setImeVisibility(true, statsToken);
                         }
                     } else {
                         showCurrentInputLocked(mLastImeTargetWindow, statsToken, flags,
@@ -5074,7 +5081,8 @@
                                 && userData.mImeBindingState.mFocusedWindowClient != null
                                 && userData.mImeBindingState.mFocusedWindowClient.mClient != null) {
                             userData.mImeBindingState.mFocusedWindowClient.mClient
-                                    .setImeVisibility(false);
+                                    .setImeVisibility(false,
+                                    null /* TODO(b329229469) check statsToken */);
                         }
                     } else {
                         @SoftInputShowHideReason final int reason = (int) msg.obj;
@@ -6822,7 +6830,8 @@
                                     && userData.mImeBindingState.mFocusedWindowClient.mClient
                                     != null) {
                                 userData.mImeBindingState.mFocusedWindowClient.mClient
-                                        .setImeVisibility(false);
+                                        .setImeVisibility(false,
+                                        null /* TODO(b329229469) initialize statsToken here? */);
                             } else {
                                 // TODO(b329229469): ImeTracker?
                             }
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index f22e654..1efb3ef 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -7180,8 +7180,8 @@
         @Override
         public void setImeInputTargetRequestedVisibility(boolean visible) {
             if (android.view.inputmethod.Flags.refactorInsetsController()) {
+                // TODO(b/329229469) we won't have the statsToken in all cases, but should still log
                 try {
-                    // TODO stats token
                     mRemoteInsetsController.setImeInputTargetRequestedVisibility(visible);
                 } catch (RemoteException e) {
                     Slog.w(TAG, "Failed to deliver setImeInputTargetRequestedVisibility", e);
diff --git a/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java b/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
index 91c61b1..4f4daa1 100644
--- a/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
+++ b/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
@@ -89,6 +89,7 @@
 
     @Override
     void onPostLayout() {
+        boolean wasServerVisible = mServerVisible;
         super.onPostLayout();
 
         if (android.view.inputmethod.Flags.refactorInsetsController()) {
@@ -101,11 +102,32 @@
             // again, so that the control with leash can be eventually dispatched
             if (!mGivenInsetsReady && mServerVisible && !givenInsetsPending) {
                 mGivenInsetsReady = true;
+                ImeTracker.forLogging().onProgress(mStatsToken,
+                        ImeTracker.PHASE_WM_POST_LAYOUT_NOTIFY_CONTROLS_CHANGED);
                 mStateController.notifyControlChanged(mControlTarget);
+            } else if (wasServerVisible && mServerVisible && mGivenInsetsReady
+                    && givenInsetsPending) {
+                // If the server visibility didn't change (still visible), and mGivenInsetsReady
+                // is set, we won't call into notifyControlChanged. Therefore, we can reset the
+                // statsToken, if available.
+                ImeTracker.forLogging().onCancelled(mStatsToken,
+                        ImeTracker.PHASE_WM_POST_LAYOUT_NOTIFY_CONTROLS_CHANGED);
+                mStatsToken = null;
             }
         }
     }
 
+    @Nullable
+    ImeTracker.Token getAndClearStatsToken() {
+        if (android.view.inputmethod.Flags.refactorInsetsController()) {
+            ImeTracker.Token statsToken = mStatsToken;
+            mStatsToken = null;
+            return statsToken;
+        } else {
+            return null;
+        }
+    }
+
     @Override
     protected boolean isLeashReadyForDispatching() {
         if (android.view.inputmethod.Flags.refactorInsetsController()) {
@@ -142,6 +164,14 @@
             }
             control.setSkipAnimationOnce(startingData != null && startingData.hasImeSurface());
         }
+        if (android.view.inputmethod.Flags.refactorInsetsController()) {
+            if (control != null && control.getLeash() != null) {
+                ImeTracker.Token statsToken = getAndClearStatsToken();
+                ImeTracker.forLogging().onProgress(statsToken,
+                        ImeTracker.PHASE_WM_GET_CONTROL_WITH_LEASH);
+                control.setImeStatsToken(statsToken);
+            }
+        }
         return control;
     }
 
@@ -168,7 +198,7 @@
             if (android.view.inputmethod.Flags.refactorInsetsController()) {
                 if (!serverVisible && !mFrozen) {
                     mGivenInsetsReady = false;
-                    updateControlForTarget(mControlTarget, true /* force */);
+                    updateControlForTarget(mControlTarget, true /* force */, null /* statsToken */);
                 }
             }
         }
@@ -214,22 +244,29 @@
     }
 
     @Override
-    void updateControlForTarget(@Nullable InsetsControlTarget target, boolean force) {
+    void updateControlForTarget(@Nullable InsetsControlTarget target, boolean force,
+            @NonNull ImeTracker.Token statsToken) {
         if (target != null && target.getWindow() != null) {
             // ime control target could be a different window.
             // Refer WindowState#getImeControlTarget().
             target = target.getWindow().getImeControlTarget();
         }
-        super.updateControlForTarget(target, force);
+        // TODO(b/329229469) make sure that the statsToken of all callers is non-null (currently
+        //  not the case)
+        super.updateControlForTarget(target, force, statsToken);
         if (Flags.refactorInsetsController()) {
+            // TODO(b/353463205) investigate if we should fail the statsToken, or if it's only
+            //  temporary null.
             if (target != null) {
-                invokeOnImeRequestedChangedListener(target.getWindow());
+                invokeOnImeRequestedChangedListener(target.getWindow(), statsToken);
             }
         }
     }
 
+    // TODO(b/353463205) change statsToken to be NonNull, after the flag is permanently enabled
     @Override
-    protected boolean updateClientVisibility(InsetsControlTarget caller) {
+    protected boolean updateClientVisibility(InsetsControlTarget caller,
+            @Nullable ImeTracker.Token statsToken) {
         InsetsControlTarget controlTarget = getControlTarget();
         if (caller != controlTarget) {
             if (Flags.refactorInsetsController()) {
@@ -240,6 +277,8 @@
                     // its new requested visibility for the IME
                     boolean imeVisible = caller.isRequestedVisible(WindowInsets.Type.ime());
                     if (controlTarget != null) {
+                        ImeTracker.forLogging().onProgress(statsToken,
+                                ImeTracker.PHASE_WM_SET_REMOTE_TARGET_IME_VISIBILITY);
                         controlTarget.setImeInputTargetRequestedVisibility(imeVisible);
                     } else {
                         // In case of a virtual display that cannot show the IME, the
@@ -249,17 +288,24 @@
                         controlTarget = mDisplayContent.getImeHostOrFallback(caller.getWindow());
 
                         if (controlTarget != caller) {
+                            ImeTracker.forLogging().onProgress(statsToken,
+                                    ImeTracker.PHASE_WM_SET_REMOTE_TARGET_IME_VISIBILITY);
                             controlTarget.setImeInputTargetRequestedVisibility(imeVisible);
+                        } else {
+                            ImeTracker.forLogging().onFailed(statsToken,
+                                    ImeTracker.PHASE_WM_SET_REMOTE_TARGET_IME_VISIBILITY);
                         }
                     }
 
                     WindowState windowState = caller.getWindow();
-                    invokeOnImeRequestedChangedListener(windowState);
+                    invokeOnImeRequestedChangedListener(windowState, statsToken);
+                } else {
+                    // TODO(b/353463205) add ImeTracker?
                 }
             }
             return false;
         }
-        boolean changed = super.updateClientVisibility(caller);
+        boolean changed = super.updateClientVisibility(caller, statsToken);
         if (!Flags.refactorInsetsController()) {
             if (changed && caller.isRequestedVisible(mSource.getType())) {
                 reportImeDrawnForOrganizerIfNeeded(caller);
@@ -273,7 +319,11 @@
                 WindowState windowState = caller.getWindow() != null ? caller.getWindow()
                         : ((mDisplayContent.getImeInputTarget() != null)
                                 ? mDisplayContent.getImeInputTarget().getWindowState() : null);
-                invokeOnImeRequestedChangedListener(windowState);
+                invokeOnImeRequestedChangedListener(windowState, statsToken);
+            } else {
+                // TODO(b/329229469) change phase and check cancelled / failed
+                ImeTracker.forLogging().onCancelled(statsToken,
+                        ImeTracker.PHASE_CLIENT_REPORT_REQUESTED_VISIBLE_TYPES);
             }
         }
         return changed;
@@ -288,22 +338,34 @@
                 // know about the new requestedVisibleTypes for the IME.
                 if (imeControlTarget != null) {
                     imeControlTarget.setImeInputTargetRequestedVisibility(
-                            (targetWin.getRequestedVisibleTypes()
-                                    & WindowInsets.Type.ime()) != 0);
+                            (targetWin.getRequestedVisibleTypes() & WindowInsets.Type.ime()) != 0);
                 }
             }
         }
     }
 
-    private void invokeOnImeRequestedChangedListener(WindowState windowState) {
+    // TODO(b/353463205) check callers to see if we can make statsToken @NonNull
+    private void invokeOnImeRequestedChangedListener(WindowState windowState,
+            @Nullable ImeTracker.Token statsToken) {
         final var imeListener = mDisplayContent.mWmService.mOnImeRequestedChangedListener;
         if (imeListener != null) {
             if (windowState != null) {
+                ImeTracker.forLogging().onProgress(statsToken,
+                        ImeTracker.PHASE_WM_POSTING_CHANGED_IME_VISIBILITY);
                 mDisplayContent.mWmService.mH.post(() -> {
+                    ImeTracker.forLogging().onProgress(statsToken,
+                            ImeTracker.PHASE_WM_INVOKING_IME_REQUESTED_LISTENER);
                     imeListener.onImeRequestedChanged(windowState.mClient.asBinder(),
-                            windowState.isRequestedVisible(WindowInsets.Type.ime()));
+                            windowState.isRequestedVisible(WindowInsets.Type.ime()), statsToken);
                 });
+            } else {
+                ImeTracker.forLogging().onFailed(statsToken,
+                        ImeTracker.PHASE_WM_POSTING_CHANGED_IME_VISIBILITY);
             }
+        } else {
+            // TODO(b/353463205) We could combine the upper if's and remove the additional phase.
+            ImeTracker.forLogging().onFailed(statsToken,
+                    ImeTracker.PHASE_WM_DISPATCH_IME_REQUESTED_CHANGED);
         }
     }
 
@@ -399,9 +461,9 @@
             // This can later become ready, so we don't want to cancel the pending request here.
             return;
         }
+        // TODO(b/329229469) check if this is still triggered, as we don't go into STATE_SHOW_IME
+        //  (DefaultImeVisibilityApplier)
         if (android.view.inputmethod.Flags.refactorInsetsController()) {
-            // Clear token here so we don't report an error in abortShowImePostLayout().
-            abortShowImePostLayout();
             // The IME is drawn, so call into {@link WindowState#notifyInsetsControlChanged}
             // if we have a leash
             if (mControl != null && mControl.getLeash() != null
@@ -517,6 +579,33 @@
     }
 
     /**
+     * Sets the statsToken before the IMS was shown/hidden.
+     * @param visible {@code true} to make it visible, false to hide it.
+     * @param statsToken the token tracking the current IME request.
+     */
+    void receiveImeStatsToken(boolean visible,
+            @NonNull ImeTracker.Token statsToken) {
+        if (!android.view.inputmethod.Flags.refactorInsetsController()) {
+            return;
+        }
+
+        if (mStatsToken != null) {
+            // We have an ongoing show request will be cancelled by the newly received show
+            // request (cancelling the initial show) or hide request (aborting the initial show).
+            logIsScheduledAndReadyToShowIme(!visible /* aborted */);
+        }
+        if (visible) {
+            ImeTracker.forLogging().onCancelled(
+                    mStatsToken, ImeTracker.PHASE_WM_ABORT_SHOW_IME_POST_LAYOUT);
+            mStatsToken = statsToken;
+        } else {
+            ImeTracker.forLogging().onFailed(
+                    mStatsToken, ImeTracker.PHASE_WM_ABORT_SHOW_IME_POST_LAYOUT);
+            mStatsToken = null;
+        }
+    }
+
+    /**
      * Logs the current state that can be checked by {@link #isScheduledAndReadyToShowIme}.
      *
      * @param aborted whether the scheduled show IME request was aborted or cancelled.
diff --git a/services/core/java/com/android/server/wm/InsetsPolicy.java b/services/core/java/com/android/server/wm/InsetsPolicy.java
index 6288a42..62bef74 100644
--- a/services/core/java/com/android/server/wm/InsetsPolicy.java
+++ b/services/core/java/com/android/server/wm/InsetsPolicy.java
@@ -47,6 +47,7 @@
 import android.view.WindowInsetsAnimation;
 import android.view.WindowInsetsAnimation.Bounds;
 import android.view.WindowManager;
+import android.view.inputmethod.ImeTracker;
 import android.view.inputmethod.InputMethodManager;
 
 import com.android.internal.R;
@@ -433,8 +434,9 @@
         return originalState;
     }
 
-    void onRequestedVisibleTypesChanged(InsetsControlTarget caller) {
-        mStateController.onRequestedVisibleTypesChanged(caller);
+    void onRequestedVisibleTypesChanged(InsetsControlTarget caller,
+            @Nullable ImeTracker.Token statsToken) {
+        mStateController.onRequestedVisibleTypesChanged(caller, statsToken);
         checkAbortTransient(caller);
         updateBarControlTarget(mFocusedWin);
     }
@@ -803,7 +805,8 @@
         }
 
         @Override
-        public void updateRequestedVisibleTypes(int types) { }
+        public void updateRequestedVisibleTypes(int types, @Nullable ImeTracker.Token statsToken) {
+        }
 
         @Override
         public boolean hasAnimationCallbacks() {
diff --git a/services/core/java/com/android/server/wm/InsetsSourceProvider.java b/services/core/java/com/android/server/wm/InsetsSourceProvider.java
index 33dea54..f5c92f6 100644
--- a/services/core/java/com/android/server/wm/InsetsSourceProvider.java
+++ b/services/core/java/com/android/server/wm/InsetsSourceProvider.java
@@ -46,6 +46,7 @@
 import android.view.SurfaceControl;
 import android.view.SurfaceControl.Transaction;
 import android.view.WindowInsets;
+import android.view.inputmethod.ImeTracker;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.protolog.ProtoLog;
@@ -498,7 +499,8 @@
         );
     }
 
-    void updateControlForTarget(@Nullable InsetsControlTarget target, boolean force) {
+    void updateControlForTarget(@Nullable InsetsControlTarget target, boolean force,
+            @Nullable ImeTracker.Token statsToken) {
         if (mSeamlessRotating) {
             // We are un-rotating the window against the display rotation. We don't want the target
             // to control the window for now.
@@ -570,7 +572,8 @@
         mSeamlessRotating = false;
     }
 
-    boolean updateClientVisibility(InsetsControlTarget caller) {
+    boolean updateClientVisibility(InsetsControlTarget caller,
+            @Nullable ImeTracker.Token statsToken) {
         final boolean requestedVisible = caller.isRequestedVisible(mSource.getType());
         if (caller != mControlTarget || requestedVisible == mClientVisible) {
             return false;
diff --git a/services/core/java/com/android/server/wm/InsetsStateController.java b/services/core/java/com/android/server/wm/InsetsStateController.java
index dcadb0f..9c2a8de 100644
--- a/services/core/java/com/android/server/wm/InsetsStateController.java
+++ b/services/core/java/com/android/server/wm/InsetsStateController.java
@@ -40,6 +40,7 @@
 import android.view.InsetsState;
 import android.view.WindowInsets;
 import android.view.WindowInsets.Type.InsetsType;
+import android.view.inputmethod.ImeTracker;
 
 import com.android.internal.protolog.ProtoLog;
 import com.android.server.inputmethod.InputMethodManagerInternal;
@@ -216,10 +217,14 @@
         }
     }
 
-    void onRequestedVisibleTypesChanged(InsetsControlTarget caller) {
+    void onRequestedVisibleTypesChanged(InsetsControlTarget caller,
+            @Nullable ImeTracker.Token statsToken) {
         boolean changed = false;
         for (int i = mProviders.size() - 1; i >= 0; i--) {
-            changed |= mProviders.valueAt(i).updateClientVisibility(caller);
+            final InsetsSourceProvider provider = mProviders.valueAt(i);
+            final boolean isImeProvider = provider.getSource().getType() == WindowInsets.Type.ime();
+            changed |= provider.updateClientVisibility(caller,
+                    isImeProvider ? statsToken : null);
         }
         if (!android.view.inputmethod.Flags.refactorInsetsController()) {
             if (changed) {
@@ -310,7 +315,9 @@
             // aborted.
             provider.updateFakeControlTarget(target);
         } else {
-            provider.updateControlForTarget(target, false /* force */);
+            // TODO(b/329229469) if the IME controlTarget changes, any pending requests should fail
+            provider.updateControlForTarget(target, false /* force */,
+                    null /* TODO(b/329229469) check if needed here */);
 
             // Get control target again in case the provider didn't accept the one we passed to it.
             target = provider.getControlTarget();
@@ -394,7 +401,9 @@
             // to the clients, so that the clients can change the current visibilities to the
             // requested visibilities with animations.
             for (int i = newControlTargets.size() - 1; i >= 0; i--) {
-                onRequestedVisibleTypesChanged(newControlTargets.valueAt(i));
+                // TODO(b/353463205) the statsToken shouldn't be null as it is used later in the
+                //  IME provider. Check if we have to create a new request here
+                onRequestedVisibleTypesChanged(newControlTargets.valueAt(i), null /* statsToken */);
             }
             newControlTargets.clear();
             if (!android.view.inputmethod.Flags.refactorInsetsController()) {
diff --git a/services/core/java/com/android/server/wm/Session.java b/services/core/java/com/android/server/wm/Session.java
index cc95518..32ec020 100644
--- a/services/core/java/com/android/server/wm/Session.java
+++ b/services/core/java/com/android/server/wm/Session.java
@@ -37,6 +37,7 @@
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_OVERLAY;
 import static android.view.WindowManager.LayoutParams.isSystemAlertWindowType;
 
+import static com.android.internal.protolog.ProtoLogGroup.WM_DEBUG_IME;
 import static com.android.internal.protolog.ProtoLogGroup.WM_SHOW_TRANSACTIONS;
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
@@ -78,6 +79,7 @@
 import android.view.WindowInsets.Type.InsetsType;
 import android.view.WindowManager;
 import android.view.WindowRelayoutResult;
+import android.view.inputmethod.ImeTracker;
 import android.window.InputTransferToken;
 import android.window.OnBackInvokedCallbackInfo;
 
@@ -710,13 +712,20 @@
     }
 
     @Override
-    public void updateRequestedVisibleTypes(IWindow window, @InsetsType int requestedVisibleTypes) {
+    public void updateRequestedVisibleTypes(IWindow window, @InsetsType int requestedVisibleTypes,
+            @Nullable ImeTracker.Token imeStatsToken) {
         synchronized (mService.mGlobalLock) {
             final WindowState win = mService.windowForClientLocked(this, window,
                     false /* throwOnError */);
             if (win != null) {
+                ImeTracker.forLogging().onProgress(imeStatsToken,
+                        ImeTracker.PHASE_WM_UPDATE_REQUESTED_VISIBLE_TYPES);
                 win.setRequestedVisibleTypes(requestedVisibleTypes);
-                win.getDisplayContent().getInsetsPolicy().onRequestedVisibleTypesChanged(win);
+                win.getDisplayContent().getInsetsPolicy().onRequestedVisibleTypesChanged(win,
+                        imeStatsToken);
+            } else {
+                ImeTracker.forLogging().onFailed(imeStatsToken,
+                        ImeTracker.PHASE_WM_UPDATE_REQUESTED_VISIBLE_TYPES);
             }
         }
     }
@@ -985,4 +994,28 @@
             }
         }
     }
+
+    @Override
+    public void notifyImeWindowVisibilityChangedFromClient(IWindow window, boolean visible,
+            @NonNull ImeTracker.Token statsToken) {
+        synchronized (mService.mGlobalLock) {
+            // TODO(b/353463205) check if we can use mService.getDefaultDisplayContentLocked()
+            //  instead of window
+            final WindowState win = mService.windowForClientLocked(this, window,
+                    false /* throwOnError */);
+            if (win != null) {
+                final InsetsStateController insetsStateController =
+                        win.getDisplayContent().getInsetsStateController();
+                ProtoLog.d(WM_DEBUG_IME, "notifyImeWindowVisibilityChangedFromClient: %s",
+                        insetsStateController.getImeSourceProvider());
+                ImeTracker.forLogging().onProgress(statsToken,
+                        ImeTracker.PHASE_WM_NOTIFY_IME_VISIBILITY_CHANGED_FROM_CLIENT);
+                insetsStateController.getImeSourceProvider().receiveImeStatsToken(visible,
+                        statsToken);
+            } else {
+                ImeTracker.forLogging().onFailed(statsToken,
+                        ImeTracker.PHASE_WM_NOTIFY_IME_VISIBILITY_CHANGED_FROM_CLIENT);
+            }
+        }
+    }
 }
diff --git a/services/core/java/com/android/server/wm/WindowManagerInternal.java b/services/core/java/com/android/server/wm/WindowManagerInternal.java
index 5061133..8ae1cf0 100644
--- a/services/core/java/com/android/server/wm/WindowManagerInternal.java
+++ b/services/core/java/com/android/server/wm/WindowManagerInternal.java
@@ -365,8 +365,10 @@
          *
          * @param windowToken The window token
          * @param imeVisible {@code true} if the IME should be shown, {@code false} to hide
+         * @param statsToken the token tracking the current IME request.
          */
-        void onImeRequestedChanged(IBinder windowToken, boolean imeVisible);
+        void onImeRequestedChanged(IBinder windowToken, boolean imeVisible,
+                @Nullable ImeTracker.Token statsToken);
     }
 
     /**
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index f65eea0..a218068 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -4717,8 +4717,10 @@
                     return;
                 }
                 dc.mRemoteInsetsControlTarget.setRequestedVisibleTypes(requestedVisibleTypes);
+                // TODO(b/353463205) the statsToken shouldn't be null as it is used later in the
+                //  IME provider. Check if we have to create a new request here
                 dc.getInsetsStateController().onRequestedVisibleTypesChanged(
-                        dc.mRemoteInsetsControlTarget);
+                        dc.mRemoteInsetsControlTarget, null /* statsToken */);
             }
         } finally {
             Binder.restoreCallingIdentity(origId);
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
index 24fc7ee..421bbae 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
@@ -3354,7 +3354,8 @@
 
         // app1 requests IME visible.
         app1.setRequestedVisibleTypes(ime(), ime());
-        mDisplayContent.getInsetsStateController().onRequestedVisibleTypesChanged(app1);
+        mDisplayContent.getInsetsStateController().onRequestedVisibleTypesChanged(app1,
+                null /* statsToken */);
 
         // Verify app1's IME insets is visible and app2's IME insets frozen flag set.
         assertTrue(app1.getInsetsState().peekSource(ID_IME).isVisible());
@@ -3425,7 +3426,7 @@
         assertFalse(activity2.mImeInsetsFrozenUntilStartInput);
 
         app1.setRequestedVisibleTypes(ime());
-        controller.onRequestedVisibleTypesChanged(app1);
+        controller.onRequestedVisibleTypesChanged(app1, null /* statsToken */);
 
         // Expect all activities in split-screen will get IME insets visible state
         assertTrue(app1.getInsetsState().peekSource(ID_IME).isVisible());
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyTests.java b/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyTests.java
index c77a4d6..caeb41c 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyTests.java
@@ -554,7 +554,7 @@
         final DisplayPolicy displayPolicy = mDisplayContent.getDisplayPolicy();
         displayPolicy.addWindowLw(mNavBarWindow, mNavBarWindow.mAttrs);
         final InsetsSourceProvider navBarProvider = mNavBarWindow.getControllableInsetProvider();
-        navBarProvider.updateControlForTarget(mAppWindow, false);
+        navBarProvider.updateControlForTarget(mAppWindow, false, null /* statsToken */);
         navBarProvider.getSource().setVisible(false);
 
         displayPolicy.setCanSystemBarsBeShownByUser(false);
@@ -579,7 +579,7 @@
         final DisplayPolicy displayPolicy = mDisplayContent.getDisplayPolicy();
         displayPolicy.addWindowLw(mNavBarWindow, mNavBarWindow.mAttrs);
         final InsetsSourceProvider navBarProvider = mNavBarWindow.getControllableInsetProvider();
-        navBarProvider.updateControlForTarget(win, false);
+        navBarProvider.updateControlForTarget(win, false, null /* statsToken */);
         navBarProvider.getSource().setVisible(false);
 
         displayPolicy.setCanSystemBarsBeShownByUser(true);
diff --git a/services/tests/wmtests/src/com/android/server/wm/InsetsPolicyTest.java b/services/tests/wmtests/src/com/android/server/wm/InsetsPolicyTest.java
index eeec54f..b26c267 100644
--- a/services/tests/wmtests/src/com/android/server/wm/InsetsPolicyTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/InsetsPolicyTest.java
@@ -361,7 +361,7 @@
 
         mAppWindow.setRequestedVisibleTypes(
                 navigationBars() | statusBars(), navigationBars() | statusBars());
-        policy.onRequestedVisibleTypesChanged(mAppWindow);
+        policy.onRequestedVisibleTypesChanged(mAppWindow, null /* statsToken */);
         waitUntilWindowAnimatorIdle();
 
         controls = mDisplayContent.getInsetsStateController().getControlsForDispatch(mAppWindow);
diff --git a/services/tests/wmtests/src/com/android/server/wm/InsetsSourceProviderTest.java b/services/tests/wmtests/src/com/android/server/wm/InsetsSourceProviderTest.java
index 2a025cd..6190807 100644
--- a/services/tests/wmtests/src/com/android/server/wm/InsetsSourceProviderTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/InsetsSourceProviderTest.java
@@ -112,13 +112,13 @@
         statusBar.getFrame().set(0, 0, 500, 100);
 
         // We must not have control or control target before we have the insets source window.
-        mProvider.updateControlForTarget(target, true /* force */);
+        mProvider.updateControlForTarget(target, true /* force */, null /* statsToken */);
         assertNull(mProvider.getControl(target));
         assertNull(mProvider.getControlTarget());
 
         // We can have the control or the control target after we have the insets source window.
         mProvider.setWindowContainer(statusBar, null, null);
-        mProvider.updateControlForTarget(target, false /* force */);
+        mProvider.updateControlForTarget(target, false /* force */, null /* statsToken */);
         assertNotNull(mProvider.getControl(target));
         assertNotNull(mProvider.getControlTarget());
 
@@ -127,25 +127,25 @@
         mProvider.startSeamlessRotation();
         assertNull(mProvider.getControl(target));
         assertNull(mProvider.getControlTarget());
-        mProvider.updateControlForTarget(target, true /* force */);
+        mProvider.updateControlForTarget(target, true /* force */, null /* statsToken */);
         assertNull(mProvider.getControl(target));
         assertNull(mProvider.getControlTarget());
 
         // We can have the control and the control target after seamless rotation.
         mProvider.finishSeamlessRotation();
-        mProvider.updateControlForTarget(target, false /* force */);
+        mProvider.updateControlForTarget(target, false /* force */, null /* statsToken */);
         assertNotNull(mProvider.getControl(target));
         assertNotNull(mProvider.getControlTarget());
 
         // We can clear the control and the control target.
-        mProvider.updateControlForTarget(null, false /* force */);
+        mProvider.updateControlForTarget(null, false /* force */, null /* statsToken */);
         assertNull(mProvider.getControl(target));
         assertNull(mProvider.getControlTarget());
 
         // We must not have control or control target if the insets source window doesn't have a
         // surface.
         statusBar.setSurfaceControl(null);
-        mProvider.updateControlForTarget(target, true /* force */);
+        mProvider.updateControlForTarget(target, true /* force */, null /* statsToken */);
         assertNull(mProvider.getControl(target));
         assertNull(mProvider.getControlTarget());
     }
@@ -173,7 +173,7 @@
 
         // We must not have control or control target before we have the insets source window,
         // so also no leash.
-        mProvider.updateControlForTarget(target, true /* force */);
+        mProvider.updateControlForTarget(target, true /* force */, null /* statsToken */);
         assertNull(mProvider.getControl(target));
         assertNull(mProvider.getControlTarget());
         assertNull(mProvider.getLeash(target));
@@ -181,7 +181,7 @@
         // We can have the control or the control target after we have the insets source window,
         // but no leash as this is not yet ready for dispatching.
         mProvider.setWindowContainer(statusBar, null, null);
-        mProvider.updateControlForTarget(target, false /* force */);
+        mProvider.updateControlForTarget(target, false /* force */, null /* statsToken */);
         assertNotNull(mProvider.getControl(target));
         assertNotNull(mProvider.getControlTarget());
         assertEquals(mProvider.getControlTarget(), target);
@@ -265,9 +265,9 @@
         final WindowState target = createWindow(null, TYPE_APPLICATION, "target");
         statusBar.getFrame().set(0, 0, 500, 100);
         mProvider.setWindowContainer(statusBar, null, null);
-        mProvider.updateControlForTarget(target, false /* force */);
+        mProvider.updateControlForTarget(target, false /* force */, null /* statsToken */);
         target.setRequestedVisibleTypes(0, statusBars());
-        mProvider.updateClientVisibility(target);
+        mProvider.updateClientVisibility(target, null /* statsToken */);
         assertFalse(mSource.isVisible());
     }
 
@@ -278,7 +278,7 @@
         statusBar.getFrame().set(0, 0, 500, 100);
         mProvider.setWindowContainer(statusBar, null, null);
         target.setRequestedVisibleTypes(0, statusBars());
-        mProvider.updateClientVisibility(target);
+        mProvider.updateClientVisibility(target, null /* statsToken */);
         assertTrue(mSource.isVisible());
     }
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/InsetsStateControllerTest.java b/services/tests/wmtests/src/com/android/server/wm/InsetsStateControllerTest.java
index c69faed..0dc56f8 100644
--- a/services/tests/wmtests/src/com/android/server/wm/InsetsStateControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/InsetsStateControllerTest.java
@@ -201,7 +201,7 @@
                 .setWindowContainer(mImeWindow, null, null);
         getController().onImeControlTargetChanged(base);
         base.setRequestedVisibleTypes(ime(), ime());
-        getController().onRequestedVisibleTypesChanged(base);
+        getController().onRequestedVisibleTypesChanged(base, null /* statsToken */);
 
         // Send our spy window (app) into the system so that we can detect the invocation.
         final WindowState win = createWindow(null, TYPE_APPLICATION, "app");
@@ -445,7 +445,7 @@
 
         waitUntilHandlersIdle();
         clearInvocations(mDisplayContent);
-        imeSourceProvider.updateControlForTarget(app, false /* force */);
+        imeSourceProvider.updateControlForTarget(app, false /* force */, null /* statsToken */);
         imeSourceProvider.setClientVisible(true);
         verify(mDisplayContent).assignWindowLayers(anyBoolean());
         waitUntilHandlersIdle();
@@ -497,7 +497,7 @@
         mDisplayContent.updateImeInputAndControlTarget(app);
 
         app.setRequestedVisibleTypes(ime(), ime());
-        getController().onRequestedVisibleTypesChanged(app);
+        getController().onRequestedVisibleTypesChanged(app, null /* statsToken */);
         assertTrue(ime.getControllableInsetProvider().getSource().isVisible());
 
         getController().updateAboveInsetsState(true /* notifyInsetsChange */);
@@ -544,7 +544,7 @@
         imeInsetsProvider.setWindowContainer(mImeWindow, null, null);
         imeInsetsProvider.updateSourceFrame(mImeWindow.getFrame());
 
-        imeInsetsProvider.updateControlForTarget(app1, false);
+        imeInsetsProvider.updateControlForTarget(app1, false, null /* statsToken */);
         imeInsetsProvider.onPostLayout();
         final InsetsSourceControl control1 = imeInsetsProvider.getControl(app1);
         assertNotNull(control1);
@@ -553,7 +553,7 @@
 
         // Simulate the IME control target updated from app1 to app2 when IME insets was invisible.
         imeInsetsProvider.setServerVisible(false);
-        imeInsetsProvider.updateControlForTarget(app2, false);
+        imeInsetsProvider.updateControlForTarget(app2, false, null /* statsToken */);
 
         // Verify insetsHint of the new control is same as last IME source frame after the layout.
         imeInsetsProvider.onPostLayout();
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java
index e13376b..b46189c 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java
@@ -472,7 +472,7 @@
         app.setRequestedVisibleTypes(0, statusBars());
         mDisplayContent.getInsetsStateController()
                 .getOrCreateSourceProvider(statusBarId, statusBars())
-                .updateClientVisibility(app);
+                .updateClientVisibility(app, null /* statsToken */);
         waitUntilHandlersIdle();
         assertFalse(statusBar.isVisible());
     }