diff --git a/core/java/android/inputmethodservice/InlineSuggestionSession.java b/core/java/android/inputmethodservice/InlineSuggestionSession.java
index 1f12d2a..25e90ee 100644
--- a/core/java/android/inputmethodservice/InlineSuggestionSession.java
+++ b/core/java/android/inputmethodservice/InlineSuggestionSession.java
@@ -39,11 +39,12 @@
 import java.util.function.Supplier;
 
 /**
- * Maintains an active inline suggestion session.
+ * Maintains an active inline suggestion session with the autofill manager service.
  *
  * <p>
- * Each session corresponds to one inline suggestion request, but there may be multiple callbacks
- * with the inline suggestions response.
+ * Each session corresponds to one {@link InlineSuggestionsRequest} and one {@link
+ * IInlineSuggestionsResponseCallback}, but there may be multiple invocations of the response
+ * callback for the same field or different fields in the same component.
  */
 class InlineSuggestionSession {
 
@@ -60,6 +61,8 @@
     @NonNull
     private final Supplier<String> mClientPackageNameSupplier;
     @NonNull
+    private final Supplier<AutofillId> mClientAutofillIdSupplier;
+    @NonNull
     private final Supplier<InlineSuggestionsRequest> mRequestSupplier;
     @NonNull
     private final Supplier<IBinder> mHostInputTokenSupplier;
@@ -71,6 +74,7 @@
     InlineSuggestionSession(@NonNull ComponentName componentName,
             @NonNull IInlineSuggestionsRequestCallback callback,
             @NonNull Supplier<String> clientPackageNameSupplier,
+            @NonNull Supplier<AutofillId> clientAutofillIdSupplier,
             @NonNull Supplier<InlineSuggestionsRequest> requestSupplier,
             @NonNull Supplier<IBinder> hostInputTokenSupplier,
             @NonNull Consumer<InlineSuggestionsResponse> responseConsumer) {
@@ -78,6 +82,7 @@
         mCallback = callback;
         mResponseCallback = new InlineSuggestionsResponseCallbackImpl(this);
         mClientPackageNameSupplier = clientPackageNameSupplier;
+        mClientAutofillIdSupplier = clientAutofillIdSupplier;
         mRequestSupplier = requestSupplier;
         mHostInputTokenSupplier = hostInputTokenSupplier;
         mResponseConsumer = responseConsumer;
@@ -115,21 +120,30 @@
         }
     }
 
-    private void handleOnInlineSuggestionsResponse(@NonNull InlineSuggestionsResponse response) {
+    private void handleOnInlineSuggestionsResponse(@NonNull AutofillId fieldId,
+            @NonNull InlineSuggestionsResponse response) {
         if (mInvalidated) {
             if (DEBUG) {
                 Log.d(TAG, "handleOnInlineSuggestionsResponse() called on invalid session");
             }
             return;
         }
-        // TODO(b/149522488): checking the current focused input field to make sure we don't send
-        //  inline responses for previous input field
+        // TODO(b/149522488): Verify fieldId against {@code mClientAutofillIdSupplier.get()} using
+        //  {@link AutofillId#equalsIgnoreSession(AutofillId)}. Right now, this seems to be
+        //  falsely alarmed quite often, depending whether autofill suggestions arrive earlier
+        //  than the IMS EditorInfo updates or not.
         if (!mComponentName.getPackageName().equals(mClientPackageNameSupplier.get())) {
             if (DEBUG) {
-                Log.d(TAG, "handleOnInlineSuggestionsResponse() called on the wrong package name");
+                Log.d(TAG,
+                        "handleOnInlineSuggestionsResponse() called on the wrong package "
+                                + "name: " + mComponentName.getPackageName() + " v.s. "
+                                + mClientPackageNameSupplier.get());
             }
             return;
         }
+        if (DEBUG) {
+            Log.d(TAG, "IME receives response: " + response.getInlineSuggestions().size());
+        }
         mResponseConsumer.accept(response);
     }
 
@@ -152,7 +166,7 @@
             if (session != null) {
                 session.mHandler.sendMessage(obtainMessage(
                         InlineSuggestionSession::handleOnInlineSuggestionsResponse, session,
-                        response));
+                        fieldId, response));
             }
         }
     }
diff --git a/core/java/android/inputmethodservice/InputMethodService.java b/core/java/android/inputmethodservice/InputMethodService.java
index 27839e7..d27d138 100644
--- a/core/java/android/inputmethodservice/InputMethodService.java
+++ b/core/java/android/inputmethodservice/InputMethodService.java
@@ -73,6 +73,7 @@
 import android.view.WindowInsets;
 import android.view.WindowManager;
 import android.view.animation.AnimationUtils;
+import android.view.autofill.AutofillId;
 import android.view.inputmethod.CompletionInfo;
 import android.view.inputmethod.CursorAnchorInfo;
 import android.view.inputmethod.EditorInfo;
@@ -825,7 +826,7 @@
             mInlineSuggestionSession.invalidateSession();
         }
         mInlineSuggestionSession = new InlineSuggestionSession(requestInfo.getComponentName(),
-                callback, this::getEditorInfoPackageName,
+                callback, this::getEditorInfoPackageName, this::getEditorInfoAutofillId,
                 () -> onCreateInlineSuggestionsRequest(requestInfo.getUiExtras()),
                 this::getHostInputToken, this::onInlineSuggestionsResponse);
     }
@@ -838,6 +839,14 @@
         return null;
     }
 
+    @Nullable
+    private AutofillId getEditorInfoAutofillId() {
+        if (mInputEditorInfo != null) {
+            return mInputEditorInfo.autofillId;
+        }
+        return null;
+    }
+
     /**
      * Returns the {@link IBinder} input token from the host view root.
      */
diff --git a/core/java/android/view/inputmethod/EditorInfo.java b/core/java/android/view/inputmethod/EditorInfo.java
index 7f90d57..6815509 100644
--- a/core/java/android/view/inputmethod/EditorInfo.java
+++ b/core/java/android/view/inputmethod/EditorInfo.java
@@ -31,6 +31,7 @@
 import android.text.TextUtils;
 import android.util.Printer;
 import android.view.View;
+import android.view.autofill.AutofillId;
 
 import com.android.internal.annotations.VisibleForTesting;
 
@@ -425,6 +426,15 @@
     public String packageName;
 
     /**
+     * Autofill Id for the field that's currently on focus.
+     *
+     * <p> Marked as hide since it's only used by framework.</p>
+     * @hide
+     */
+    @NonNull
+    public AutofillId autofillId = new AutofillId(View.NO_ID);
+
+    /**
      * Identifier for the editor's field.  This is optional, and may be
      * 0.  By default it is filled in with the result of
      * {@link android.view.View#getId() View.getId()} on the View that
@@ -793,6 +803,7 @@
         pw.println(prefix + "hintText=" + hintText
                 + " label=" + label);
         pw.println(prefix + "packageName=" + packageName
+                + " autofillId=" + autofillId
                 + " fieldId=" + fieldId
                 + " fieldName=" + fieldName);
         pw.println(prefix + "extras=" + extras);
@@ -821,6 +832,7 @@
         TextUtils.writeToParcel(hintText, dest, flags);
         TextUtils.writeToParcel(label, dest, flags);
         dest.writeString(packageName);
+        autofillId.writeToParcel(dest, flags);
         dest.writeInt(fieldId);
         dest.writeString(fieldName);
         dest.writeBundle(extras);
@@ -852,6 +864,7 @@
                     res.hintText = TextUtils.CHAR_SEQUENCE_CREATOR.createFromParcel(source);
                     res.label = TextUtils.CHAR_SEQUENCE_CREATOR.createFromParcel(source);
                     res.packageName = source.readString();
+                    res.autofillId = AutofillId.CREATOR.createFromParcel(source);
                     res.fieldId = source.readInt();
                     res.fieldName = source.readString();
                     res.extras = source.readBundle();
diff --git a/core/java/android/view/inputmethod/InputMethodManager.java b/core/java/android/view/inputmethod/InputMethodManager.java
index f3aa314..aca265b 100644
--- a/core/java/android/view/inputmethod/InputMethodManager.java
+++ b/core/java/android/view/inputmethod/InputMethodManager.java
@@ -1858,6 +1858,7 @@
         // system can verify the consistency between the uid of this process and package name passed
         // from here. See comment of Context#getOpPackageName() for details.
         tba.packageName = view.getContext().getOpPackageName();
+        tba.autofillId = view.getAutofillId();
         tba.fieldId = view.getId();
         InputConnection ic = view.onCreateInputConnection(tba);
         if (DEBUG) Log.v(TAG, "Starting input: tba=" + tba + " ic=" + ic);
diff --git a/media/java/android/media/MediaRouter2.java b/media/java/android/media/MediaRouter2.java
index 68df7d2..fd24089 100644
--- a/media/java/android/media/MediaRouter2.java
+++ b/media/java/android/media/MediaRouter2.java
@@ -1071,7 +1071,7 @@
                 try {
                     mMediaRouterService.deselectRouteWithRouter2(stub, getId(), route);
                 } catch (RemoteException ex) {
-                    Log.e(TAG, "Unable to remove route from session.", ex);
+                    Log.e(TAG, "Unable to deselect route from session.", ex);
                 }
             }
         }
diff --git a/media/java/android/media/MediaRouter2Manager.java b/media/java/android/media/MediaRouter2Manager.java
index ff2c863..fb45ae1 100644
--- a/media/java/android/media/MediaRouter2Manager.java
+++ b/media/java/android/media/MediaRouter2Manager.java
@@ -44,8 +44,10 @@
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
 
 /**
+ * A class that monitors and controls media routing of other apps.
  * @hide
  */
 public class MediaRouter2Manager {
@@ -201,10 +203,24 @@
     }
 
     /**
+     * Gets the system routing session associated with no specific application.
+     */
+    @NonNull
+    public RoutingSessionInfo getSystemRoutingSession() {
+        for (RoutingSessionInfo sessionInfo : getActiveSessions()) {
+            if (sessionInfo.isSystemSession()) {
+                return sessionInfo;
+            }
+        }
+        throw new IllegalStateException("No system routing session");
+    }
+
+    /**
      * Gets routing sessions of an application with the given package name.
-     * The first element of the returned list is the system routing controller.
+     * The first element of the returned list is the system routing session.
      *
-     * @see MediaRouter2#getSystemController()
+     * @param packageName the package name of the application that is routing.
+     * @see #getSystemRoutingSession()
      */
     @NonNull
     public List<RoutingSessionInfo> getRoutingSessions(@NonNull String packageName) {
@@ -213,8 +229,11 @@
         List<RoutingSessionInfo> sessions = new ArrayList<>();
 
         for (RoutingSessionInfo sessionInfo : getActiveSessions()) {
-            if (sessionInfo.isSystemSession()
-                    || TextUtils.equals(sessionInfo.getClientPackageName(), packageName)) {
+            if (sessionInfo.isSystemSession()) {
+                sessions.add(new RoutingSessionInfo.Builder(sessionInfo)
+                        .setClientPackageName(packageName)
+                        .build());
+            } else if (TextUtils.equals(sessionInfo.getClientPackageName(), packageName)) {
                 sessions.add(sessionInfo);
             }
         }
@@ -223,10 +242,15 @@
 
     /**
      * Gets the list of all active routing sessions.
+     * <p>
      * The first element of the list is the system routing session containing
      * phone speakers, wired headset, Bluetooth devices.
      * The system routing session is shared by apps such that controlling it will affect
      * all apps.
+     * If you want to transfer media of an application, use {@link #getRoutingSessions(String)}.
+     *
+     * @see #getRoutingSessions(String)
+     * @see #getSystemRoutingSession()
      */
     @NonNull
     public List<RoutingSessionInfo> getActiveSessions() {
@@ -258,31 +282,46 @@
 
     /**
      * Selects media route for the specified package name.
-     *
-     * If the given route is {@link RoutingController#getTransferableRoutes() a transferable
-     * route} of a routing session of the application, the session will be transferred to
-     * the route. If not, a new routing session will be created.
-     *
-     * @param packageName the package name of the application that should change it's media route
-     * @param route the route to be selected.
      */
     public void selectRoute(@NonNull String packageName, @NonNull MediaRoute2Info route) {
         Objects.requireNonNull(packageName, "packageName must not be null");
         Objects.requireNonNull(route, "route must not be null");
 
-        boolean transferred = false;
-        //TODO: instead of release all controllers, add an API to specify controllers that
-        // should be released (or is the system controller).
-        for (RoutingSessionInfo sessionInfo : getRoutingSessions(packageName)) {
-            if (!transferred && sessionInfo.getTransferableRoutes().contains(route.getId())) {
-                new RoutingController(sessionInfo).transferToRoute(route);
-                transferred = true;
-            } else if (!sessionInfo.isSystemSession()) {
-                new RoutingController(sessionInfo).release();
-            }
+        List<RoutingSessionInfo> sessionInfos = getRoutingSessions(packageName);
+        RoutingSessionInfo targetSession = sessionInfos.get(sessionInfos.size() - 1);
+        transfer(targetSession, route);
+    }
+
+    /**
+     * Transfers a routing session to a media route.
+     * <p>{@link Callback#onTransferred} or {@link Callback#onTransferFailed} will be called
+     * depending on the result.
+     *
+     * @param sessionInfo the routing session info to transfer
+     * @param route the route transfer to
+     *
+     * @see Callback#onTransferred(RoutingSessionInfo, RoutingSessionInfo)
+     * @see Callback#onTransferFailed(RoutingSessionInfo, MediaRoute2Info)
+     */
+    public void transfer(@NonNull RoutingSessionInfo sessionInfo,
+            @Nullable MediaRoute2Info route) {
+        Objects.requireNonNull(sessionInfo, "sessionInfo must not be null");
+
+        if (route == null) {
+            releaseSession(sessionInfo);
+            return;
         }
 
-        if (transferred) {
+        //TODO: Ignore unknown route.
+        if (sessionInfo.getTransferableRoutes().contains(route.getId())) {
+            //TODO: callbacks must be called after this.
+            transferToRoute(sessionInfo, route);
+            return;
+        }
+
+        if (TextUtils.isEmpty(sessionInfo.getClientPackageName())) {
+            Log.w(TAG, "transfer: Ignoring transfer without package name.");
+            notifyTransferFailed(sessionInfo, route);
             return;
         }
 
@@ -294,7 +333,7 @@
             try {
                 int requestId = mNextRequestId.getAndIncrement();
                 mMediaRouterService.requestCreateSessionWithManager(
-                        client, packageName, route, requestId);
+                        client, sessionInfo.getClientPackageName(), route, requestId);
                 //TODO: release the previous session?
             } catch (RemoteException ex) {
                 Log.e(TAG, "Unable to select media route", ex);
@@ -451,6 +490,18 @@
         }
     }
 
+    void notifyTransferred(RoutingSessionInfo oldSession, RoutingSessionInfo newSession) {
+        for (CallbackRecord record : mCallbackRecords) {
+            record.mExecutor.execute(() -> record.mCallback.onTransferred(oldSession, newSession));
+        }
+    }
+
+    void notifyTransferFailed(RoutingSessionInfo sessionInfo, MediaRoute2Info route) {
+        for (CallbackRecord record : mCallbackRecords) {
+            record.mExecutor.execute(() -> record.mCallback.onTransferFailed(sessionInfo, route));
+        }
+    }
+
     void updatePreferredFeatures(String packageName, List<String> preferredFeatures) {
         List<String> prevFeatures = mPreferredFeaturesMap.put(packageName, preferredFeatures);
         if ((prevFeatures == null && preferredFeatures.size() == 0)
@@ -475,6 +526,204 @@
     }
 
     /**
+     * Gets the unmodifiable list of selected routes for the session.
+     */
+    @NonNull
+    public List<MediaRoute2Info> getSelectedRoutes(@NonNull RoutingSessionInfo sessionInfo) {
+        Objects.requireNonNull(sessionInfo, "sessionInfo must not be null");
+
+        List<String> routeIds = sessionInfo.getSelectedRoutes();
+        return getRoutesWithIds(routeIds);
+    }
+
+    /**
+     * Gets the unmodifiable list of selectable routes for the session.
+     */
+    @NonNull
+    public List<MediaRoute2Info> getSelectableRoutes(@NonNull RoutingSessionInfo sessionInfo) {
+        Objects.requireNonNull(sessionInfo, "sessionInfo must not be null");
+
+        List<String> routeIds = sessionInfo.getSelectableRoutes();
+        return getRoutesWithIds(routeIds);
+    }
+
+    /**
+     * Gets the unmodifiable list of deselectable routes for the session.
+     */
+    @NonNull
+    public List<MediaRoute2Info> getDeselectableRoutes(@NonNull RoutingSessionInfo sessionInfo) {
+        Objects.requireNonNull(sessionInfo, "sessionInfo must not be null");
+
+        List<String> routeIds = sessionInfo.getDeselectableRoutes();
+        return getRoutesWithIds(routeIds);
+    }
+
+    /**
+     * Selects a route for the remote session. After a route is selected, the media is expected
+     * to be played to the all the selected routes. This is different from {@link
+     * #transfer(RoutingSessionInfo, MediaRoute2Info)} transferring to a route},
+     * where the media is expected to 'move' from one route to another.
+     * <p>
+     * The given route must satisfy all of the following conditions:
+     * <ul>
+     * <li>it should not be included in {@link #getSelectedRoutes(RoutingSessionInfo)}</li>
+     * <li>it should be included in {@link #getSelectableRoutes(RoutingSessionInfo)}</li>
+     * </ul>
+     * If the route doesn't meet any of above conditions, it will be ignored.
+     *
+     * @see #getSelectedRoutes(RoutingSessionInfo)
+     * @see #getSelectableRoutes(RoutingSessionInfo)
+     * @see Callback#onSessionsUpdated()
+     */
+    public void selectRoute(@NonNull RoutingSessionInfo sessionInfo,
+            @NonNull MediaRoute2Info route) {
+        Objects.requireNonNull(sessionInfo, "sessionInfo must not be null");
+        Objects.requireNonNull(route, "route must not be null");
+
+        if (sessionInfo.getSelectedRoutes().contains(route.getId())) {
+            Log.w(TAG, "Ignoring selecting a route that is already selected. route=" + route);
+            return;
+        }
+
+        if (!sessionInfo.getSelectableRoutes().contains(route.getId())) {
+            Log.w(TAG, "Ignoring selecting a non-selectable route=" + route);
+            return;
+        }
+
+        Client client;
+        synchronized (sLock) {
+            client = mClient;
+        }
+        if (client != null) {
+            try {
+                int requestId = mNextRequestId.getAndIncrement();
+                mMediaRouterService.selectRouteWithManager(
+                        mClient, sessionInfo.getId(), route, requestId);
+            } catch (RemoteException ex) {
+                Log.e(TAG, "selectRoute: Failed to send a request.", ex);
+            }
+        }
+    }
+
+    /**
+     * Deselects a route from the remote session. After a route is deselected, the media is
+     * expected to be stopped on the deselected routes.
+     * <p>
+     * The given route must satisfy all of the following conditions:
+     * <ul>
+     * <li>it should be included in {@link #getSelectedRoutes(RoutingSessionInfo)}</li>
+     * <li>it should be included in {@link #getDeselectableRoutes(RoutingSessionInfo)}</li>
+     * </ul>
+     * If the route doesn't meet any of above conditions, it will be ignored.
+     *
+     * @see #getSelectedRoutes(RoutingSessionInfo)
+     * @see #getDeselectableRoutes(RoutingSessionInfo)
+     * @see Callback#onSessionsUpdated()
+     */
+    public void deselectRoute(@NonNull RoutingSessionInfo sessionInfo,
+            @NonNull MediaRoute2Info route) {
+        Objects.requireNonNull(sessionInfo, "sessionInfo must not be null");
+        Objects.requireNonNull(route, "route must not be null");
+
+        if (!sessionInfo.getSelectedRoutes().contains(route.getId())) {
+            Log.w(TAG, "Ignoring deselecting a route that is not selected. route=" + route);
+            return;
+        }
+
+        if (!sessionInfo.getDeselectableRoutes().contains(route.getId())) {
+            Log.w(TAG, "Ignoring deselecting a non-deselectable route=" + route);
+            return;
+        }
+
+        Client client;
+        synchronized (sLock) {
+            client = mClient;
+        }
+        if (client != null) {
+            try {
+                int requestId = mNextRequestId.getAndIncrement();
+                mMediaRouterService.deselectRouteWithManager(
+                        mClient, sessionInfo.getId(), route, requestId);
+            } catch (RemoteException ex) {
+                Log.e(TAG, "deselectRoute: Failed to send a request.", ex);
+            }
+        }
+    }
+
+    /**
+     * Transfers to a given route for the remote session.
+     *
+     * @hide
+     */
+    void transferToRoute(@NonNull RoutingSessionInfo sessionInfo,
+            @NonNull MediaRoute2Info route) {
+        Objects.requireNonNull(sessionInfo, "sessionInfo must not be null");
+        Objects.requireNonNull(route, "route must not be null");
+
+        if (sessionInfo.getSelectedRoutes().contains(route.getId())) {
+            Log.w(TAG, "Ignoring transferring to a route that is already added. route="
+                    + route);
+            return;
+        }
+
+        if (!sessionInfo.getTransferableRoutes().contains(route.getId())) {
+            Log.w(TAG, "Ignoring transferring to a non-transferable route=" + route);
+            return;
+        }
+
+        Client client;
+        synchronized (sLock) {
+            client = mClient;
+        }
+        if (client != null) {
+            try {
+                int requestId = mNextRequestId.getAndIncrement();
+                mMediaRouterService.transferToRouteWithManager(
+                        mClient, sessionInfo.getId(), route, requestId);
+            } catch (RemoteException ex) {
+                Log.e(TAG, "transferToRoute: Failed to send a request.", ex);
+            }
+        }
+    }
+
+    /**
+     * Requests releasing a session.
+     * <p>
+     * If a session is released, any operation on the session will be ignored.
+     * {@link Callback#onTransferred(RoutingSessionInfo, RoutingSessionInfo)} with {@code null}
+     * session will be called when the session is released.
+     * </p>
+     *
+     * @see Callback#onTransferred(RoutingSessionInfo, RoutingSessionInfo)
+     */
+    public void releaseSession(@NonNull RoutingSessionInfo sessionInfo) {
+        Objects.requireNonNull(sessionInfo, "sessionInfo must not be null");
+
+        Client client;
+        synchronized (sLock) {
+            client = mClient;
+        }
+        if (client != null) {
+            try {
+                int requestId = mNextRequestId.getAndIncrement();
+                mMediaRouterService.releaseSessionWithManager(
+                        mClient, sessionInfo.getId(), requestId);
+            } catch (RemoteException ex) {
+                Log.e(TAG, "releaseSession: Failed to send a request", ex);
+            }
+        }
+    }
+
+    private List<MediaRoute2Info> getRoutesWithIds(List<String> routeIds) {
+        synchronized (sLock) {
+            return routeIds.stream().map(mRoutes::get)
+                    .filter(Objects::nonNull)
+                    .collect(Collectors.toList());
+        }
+    }
+
+    //TODO: Remove this.
+    /**
      * A class to control media routing session in media route provider.
      * With routing controller, an application can select a route into the session or deselect
      * a route in the session.
@@ -489,6 +738,15 @@
         }
 
         /**
+         * Releases the session
+         */
+        public void release() {
+            synchronized (mControllerLock) {
+                releaseSession(mSessionInfo);
+            }
+        }
+
+        /**
          * Gets the ID of the session
          */
         @NonNull
@@ -523,11 +781,7 @@
          */
         @NonNull
         public List<MediaRoute2Info> getSelectedRoutes() {
-            List<String> routeIds;
-            synchronized (mControllerLock) {
-                routeIds = mSessionInfo.getSelectedRoutes();
-            }
-            return getRoutesWithIds(routeIds);
+            return MediaRouter2Manager.this.getSelectedRoutes(mSessionInfo);
         }
 
         /**
@@ -535,11 +789,7 @@
          */
         @NonNull
         public List<MediaRoute2Info> getSelectableRoutes() {
-            List<String> routeIds;
-            synchronized (mControllerLock) {
-                routeIds = mSessionInfo.getSelectableRoutes();
-            }
-            return getRoutesWithIds(routeIds);
+            return MediaRouter2Manager.this.getSelectableRoutes(mSessionInfo);
         }
 
         /**
@@ -547,11 +797,7 @@
          */
         @NonNull
         public List<MediaRoute2Info> getDeselectableRoutes() {
-            List<String> routeIds;
-            synchronized (mControllerLock) {
-                routeIds = mSessionInfo.getDeselectableRoutes();
-            }
-            return getRoutesWithIds(routeIds);
+            return MediaRouter2Manager.this.getDeselectableRoutes(mSessionInfo);
         }
 
         /**
@@ -579,35 +825,7 @@
          * @see #getSelectableRoutes()
          */
         public void selectRoute(@NonNull MediaRoute2Info route) {
-            Objects.requireNonNull(route, "route must not be null");
-
-            RoutingSessionInfo sessionInfo;
-            synchronized (mControllerLock) {
-                sessionInfo = mSessionInfo;
-            }
-            if (sessionInfo.getSelectedRoutes().contains(route.getId())) {
-                Log.w(TAG, "Ignoring selecting a route that is already selected. route=" + route);
-                return;
-            }
-
-            if (!sessionInfo.getSelectableRoutes().contains(route.getId())) {
-                Log.w(TAG, "Ignoring selecting a non-selectable route=" + route);
-                return;
-            }
-
-            Client client;
-            synchronized (sLock) {
-                client = mClient;
-            }
-            if (client != null) {
-                try {
-                    int requestId = mNextRequestId.getAndIncrement();
-                    mMediaRouterService.selectRouteWithManager(
-                            mClient, getSessionId(), route, requestId);
-                } catch (RemoteException ex) {
-                    Log.e(TAG, "Unable to select route for session.", ex);
-                }
-            }
+            MediaRouter2Manager.this.selectRoute(mSessionInfo, route);
         }
 
         /**
@@ -623,104 +841,19 @@
          * @see #getDeselectableRoutes()
          */
         public void deselectRoute(@NonNull MediaRoute2Info route) {
-            Objects.requireNonNull(route, "route must not be null");
-            RoutingSessionInfo sessionInfo;
-            synchronized (mControllerLock) {
-                sessionInfo = mSessionInfo;
-            }
-
-            if (!sessionInfo.getSelectedRoutes().contains(route.getId())) {
-                Log.w(TAG, "Ignoring deselecting a route that is not selected. route=" + route);
-                return;
-            }
-
-            if (!sessionInfo.getDeselectableRoutes().contains(route.getId())) {
-                Log.w(TAG, "Ignoring deselecting a non-deselectable route=" + route);
-                return;
-            }
-
-            Client client;
-            synchronized (sLock) {
-                client = mClient;
-            }
-            if (client != null) {
-                try {
-                    int requestId = mNextRequestId.getAndIncrement();
-                    mMediaRouterService.deselectRouteWithManager(
-                            mClient, getSessionId(), route, requestId);
-                } catch (RemoteException ex) {
-                    Log.e(TAG, "Unable to remove route from session.", ex);
-                }
-            }
+            MediaRouter2Manager.this.deselectRoute(mSessionInfo, route);
         }
 
         /**
-         * Transfers to a given route for the remote session. The given route must satisfy
-         * all of the following conditions:
-         * <ul>
-         * <li>ID should not be included in {@link #getSelectedRoutes()}</li>
-         * <li>ID should be included in {@link #getTransferableRoutes()}</li>
-         * </ul>
-         * If the route doesn't meet any of above conditions, it will be ignored.
-         *
-         * @see #getSelectedRoutes()
-         * @see #getTransferableRoutes()
+         * Transfers session to the given rotue.
          */
         public void transferToRoute(@NonNull MediaRoute2Info route) {
-            Objects.requireNonNull(route, "route must not be null");
-            RoutingSessionInfo sessionInfo;
-            synchronized (mControllerLock) {
-                sessionInfo = mSessionInfo;
-            }
-
-            if (sessionInfo.getSelectedRoutes().contains(route.getId())) {
-                Log.w(TAG, "Ignoring transferring to a route that is already added. route="
-                        + route);
-                return;
-            }
-
-            if (!sessionInfo.getTransferableRoutes().contains(route.getId())) {
-                Log.w(TAG, "Ignoring transferring to a non-transferable route=" + route);
-                return;
-            }
-
-            Client client;
-            synchronized (sLock) {
-                client = mClient;
-            }
-            if (client != null) {
-                try {
-                    int requestId = mNextRequestId.getAndIncrement();
-                    mMediaRouterService.transferToRouteWithManager(
-                            mClient, getSessionId(), route, requestId);
-                } catch (RemoteException ex) {
-                    Log.e(TAG, "Unable to transfer to route for session.", ex);
-                }
-            }
-        }
-
-        /**
-         * Release this session.
-         * Any operation on this session after calling this method will be ignored.
-         */
-        public void release() {
-            Client client;
-            synchronized (sLock) {
-                client = mClient;
-            }
-            if (client != null) {
-                try {
-                    int requestId = mNextRequestId.getAndIncrement();
-                    mMediaRouterService.releaseSessionWithManager(
-                            mClient, getSessionId(), requestId);
-                } catch (RemoteException ex) {
-                    Log.e(TAG, "Unable to notify of controller release", ex);
-                }
-            }
+            MediaRouter2Manager.this.transferToRoute(mSessionInfo, route);
         }
 
         /**
          * Gets the session info of the session
+         *
          * @hide
          */
         @NonNull
@@ -729,19 +862,6 @@
                 return mSessionInfo;
             }
         }
-
-        private List<MediaRoute2Info> getRoutesWithIds(List<String> routeIds) {
-            List<MediaRoute2Info> routes = new ArrayList<>();
-            synchronized (mRoutesLock) {
-                for (String routeId : routeIds) {
-                    MediaRoute2Info route = mRoutes.get(routeId);
-                    if (route != null) {
-                        routes.add(route);
-                    }
-                }
-            }
-            return Collections.unmodifiableList(routes);
-        }
     }
 
     /**
@@ -780,7 +900,24 @@
          */
         public void onSessionsUpdated() {}
 
-        //TODO: remove this
+        //TODO: Call this.
+        /**
+         * Called when media is transferred.
+         *
+         * @param oldSession the previous session
+         * @param newSession the new session or {@code null} if the session is released.
+         */
+        public void onTransferred(@NonNull RoutingSessionInfo oldSession,
+                @Nullable RoutingSessionInfo newSession) { }
+
+        //TODO: Call this.
+        /**
+         * Called when {@link #transfer(RoutingSessionInfo, MediaRoute2Info)} fails.
+         */
+        public void onTransferFailed(@NonNull RoutingSessionInfo session,
+                @NonNull MediaRoute2Info route) { }
+
+        //TODO: Remove this.
         /**
          * Called when the preferred route features of an app is changed.
          *
diff --git a/media/java/android/media/RoutingSessionInfo.java b/media/java/android/media/RoutingSessionInfo.java
index 9d81fbb..2276b6a 100644
--- a/media/java/android/media/RoutingSessionInfo.java
+++ b/media/java/android/media/RoutingSessionInfo.java
@@ -357,7 +357,7 @@
         // TODO: Reorder these (important ones first)
         final String mId;
         CharSequence mName;
-        final String mClientPackageName;
+        String mClientPackageName;
         String mProviderId;
         final List<String> mSelectedRoutes;
         final List<String> mSelectableRoutes;
@@ -434,6 +434,17 @@
         }
 
         /**
+         * Sets the client package name of the session.
+         *
+         * @hide
+         */
+        @NonNull
+        public Builder setClientPackageName(@Nullable String packageName) {
+            mClientPackageName = packageName;
+            return this;
+        }
+
+        /**
          * Sets the provider ID of the session.
          *
          * @hide
diff --git a/media/java/android/media/session/MediaSession.java b/media/java/android/media/session/MediaSession.java
index 486c0c2..8bf462c 100644
--- a/media/java/android/media/session/MediaSession.java
+++ b/media/java/android/media/session/MediaSession.java
@@ -196,7 +196,7 @@
                 .getSystemService(Context.MEDIA_SESSION_SERVICE);
         try {
             mBinder = manager.createSession(mCbStub, tag, sessionInfo);
-            mSessionToken = new Token(mBinder.getController());
+            mSessionToken = new Token(Process.myUid(), mBinder.getController());
             mController = new MediaController(context, mSessionToken);
         } catch (RemoteException e) {
             throw new RuntimeException("Remote error creating session.", e);
@@ -771,8 +771,8 @@
         /**
          * @hide
          */
-        public Token(ISessionController binder) {
-            mUid = Process.myUid();
+        public Token(int uid, ISessionController binder) {
+            mUid = uid;
             mBinder = binder;
         }
 
diff --git a/media/tests/MediaRouter/src/com/android/mediaroutertest/MediaRouter2ManagerTest.java b/media/tests/MediaRouter/src/com/android/mediaroutertest/MediaRouter2ManagerTest.java
index 4e12859..230b9e4 100644
--- a/media/tests/MediaRouter/src/com/android/mediaroutertest/MediaRouter2ManagerTest.java
+++ b/media/tests/MediaRouter/src/com/android/mediaroutertest/MediaRouter2ManagerTest.java
@@ -248,10 +248,9 @@
 
         assertEquals(2, sessions.size());
 
-        MediaRouter2Manager.RoutingController routingController =
-                mManager.getControllerForSession(sessions.get(1));
+        RoutingSessionInfo sessionInfo = sessions.get(1);
         awaitOnRouteChangedManager(
-                () -> routingController.release(),
+                () -> mManager.releaseSession(sessionInfo),
                 ROUTE_ID1,
                 route -> TextUtils.equals(route.getClientPackageName(), null));
         assertEquals(1, mManager.getRoutingSessions(mPackageName).size());
@@ -283,8 +282,7 @@
         List<RoutingSessionInfo> sessions = mManager.getRoutingSessions(mPackageName);
 
         assertEquals(2, sessions.size());
-        MediaRouter2Manager.RoutingController routingController =
-                mManager.getControllerForSession(sessions.get(1));
+        RoutingSessionInfo sessionInfo = sessions.get(1);
 
         awaitOnRouteChangedManager(
                 () -> mManager.selectRoute(mPackageName, routes.get(ROUTE_ID5_TO_TRANSFER_TO)),
@@ -292,7 +290,7 @@
                 route -> TextUtils.equals(route.getClientPackageName(), mPackageName));
 
         awaitOnRouteChangedManager(
-                () -> routingController.release(),
+                () -> mManager.releaseSession(sessionInfo),
                 ROUTE_ID5_TO_TRANSFER_TO,
                 route -> TextUtils.equals(route.getClientPackageName(), null));
     }
@@ -573,7 +571,7 @@
         addManagerCallback(new MediaRouter2Manager.Callback());
 
         for (RoutingSessionInfo session : mManager.getActiveSessions()) {
-            mManager.getControllerForSession(session).release();
+            mManager.releaseSession(session);
         }
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogControllerImplTest.java b/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogControllerImplTest.java
index 80aa6f6..8cc83dd 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogControllerImplTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogControllerImplTest.java
@@ -29,6 +29,7 @@
 import android.media.AudioManager;
 import android.media.session.MediaSession;
 import android.os.Handler;
+import android.os.Process;
 import android.testing.AndroidTestingRunner;
 
 import androidx.test.filters.SmallTest;
@@ -116,13 +117,13 @@
 
     @Test
     public void testOnRemoteVolumeChanged_newStream_noNullPointer() {
-        MediaSession.Token token = new MediaSession.Token(null);
+        MediaSession.Token token = new MediaSession.Token(Process.myUid(), null);
         mVolumeController.mMediaSessionsCallbacksW.onRemoteVolumeChanged(token, 0);
     }
 
     @Test
     public void testOnRemoteRemove_newStream_noNullPointer() {
-        MediaSession.Token token = new MediaSession.Token(null);
+        MediaSession.Token token = new MediaSession.Token(Process.myUid(), null);
         mVolumeController.mMediaSessionsCallbacksW.onRemoteRemoved(token);
     }
 
diff --git a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-ar/strings.xml b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-ar/strings.xml
new file mode 100644
index 0000000..307bf70
--- /dev/null
+++ b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-ar/strings.xml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!-- 
+  ~ Copyright (C) 2020 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+   -->
+
+<resources xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
+    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"حواف منحنية"</string>
+</resources>
diff --git a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-bn/strings.xml b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-bn/strings.xml
new file mode 100644
index 0000000..565f75e
--- /dev/null
+++ b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-bn/strings.xml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!-- 
+  ~ Copyright (C) 2020 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+   -->
+
+<resources xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
+    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"ওয়াটারফল কাট-আউট"</string>
+</resources>
diff --git a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-bs/strings.xml b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-bs/strings.xml
index c9cff08..594999f 100644
--- a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-bs/strings.xml
+++ b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-bs/strings.xml
@@ -17,5 +17,5 @@
 
 <resources xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
-    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"Urez vodopada"</string>
+    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"Izrezani vodopad"</string>
 </resources>
diff --git a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-eu/strings.xml b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-eu/strings.xml
new file mode 100644
index 0000000..93ef2c8
--- /dev/null
+++ b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-eu/strings.xml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!-- 
+  ~ Copyright (C) 2020 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+   -->
+
+<resources xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
+    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"Ur-jauzi moduko mozketa"</string>
+</resources>
diff --git a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-fr-rCA/strings.xml b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-fr-rCA/strings.xml
new file mode 100644
index 0000000..7232c33
--- /dev/null
+++ b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-fr-rCA/strings.xml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!-- 
+  ~ Copyright (C) 2020 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+   -->
+
+<resources xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
+    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"Encoche en cascade"</string>
+</resources>
diff --git a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-gu/strings.xml b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-gu/strings.xml
new file mode 100644
index 0000000..03672fe
--- /dev/null
+++ b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-gu/strings.xml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!-- 
+  ~ Copyright (C) 2020 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+   -->
+
+<resources xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
+    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"વૉટરફૉલ કટઆઉટ"</string>
+</resources>
diff --git a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-hi/strings.xml b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-hi/strings.xml
new file mode 100644
index 0000000..319d81a
--- /dev/null
+++ b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-hi/strings.xml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!-- 
+  ~ Copyright (C) 2020 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+   -->
+
+<resources xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
+    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"वॉटरफ़ॉल कटआउट"</string>
+</resources>
diff --git a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-it/strings.xml b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-it/strings.xml
index dde9914..3ea14c5 100644
--- a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-it/strings.xml
+++ b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-it/strings.xml
@@ -17,5 +17,5 @@
 
 <resources xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
-    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"Ritaglio a cascata"</string>
+    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"Curvatura Waterfall"</string>
 </resources>
diff --git a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-iw/strings.xml b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-iw/strings.xml
index f71a879..5dbce7e 100644
--- a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-iw/strings.xml
+++ b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-iw/strings.xml
@@ -17,5 +17,5 @@
 
 <resources xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
-    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"מגרעת מפל"</string>
+    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"שוליים מעוגלים"</string>
 </resources>
diff --git a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-ja/strings.xml b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-ja/strings.xml
index 354ce59..4db0149 100644
--- a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-ja/strings.xml
+++ b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-ja/strings.xml
@@ -17,5 +17,5 @@
 
 <resources xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
-    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"ウォーターフォールのカットアウト"</string>
+    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"エッジ スクリーン"</string>
 </resources>
diff --git a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-kk/strings.xml b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-kk/strings.xml
new file mode 100644
index 0000000..bb0dfe9
--- /dev/null
+++ b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-kk/strings.xml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!-- 
+  ~ Copyright (C) 2020 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+   -->
+
+<resources xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
+    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"Сарқырама ойығы"</string>
+</resources>
diff --git a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-km/strings.xml b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-km/strings.xml
index 8d2f887..b73ccbb 100644
--- a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-km/strings.xml
+++ b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-km/strings.xml
@@ -17,5 +17,5 @@
 
 <resources xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
-    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"អេក្រង់​គ្មានគែម"</string>
+    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"គែមកោង"</string>
 </resources>
diff --git a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-ky/strings.xml b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-ky/strings.xml
new file mode 100644
index 0000000..18e2083
--- /dev/null
+++ b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-ky/strings.xml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!-- 
+  ~ Copyright (C) 2020 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+   -->
+
+<resources xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
+    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"Шаркыратманы кесүү"</string>
+</resources>
diff --git a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-mk/strings.xml b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-mk/strings.xml
index f39584b..a330a35 100644
--- a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-mk/strings.xml
+++ b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-mk/strings.xml
@@ -17,5 +17,5 @@
 
 <resources xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
-    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"Исечок во вид на водопад"</string>
+    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"Прекин за Waterfall"</string>
 </resources>
diff --git a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-ml/strings.xml b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-ml/strings.xml
new file mode 100644
index 0000000..112562d
--- /dev/null
+++ b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-ml/strings.xml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!-- 
+  ~ Copyright (C) 2020 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+   -->
+
+<resources xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
+    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"വെള്ളച്ചാട്ട കട്ടൗട്ട്"</string>
+</resources>
diff --git a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-pt-rBR/strings.xml b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-pt-rBR/strings.xml
index f80fcad..3ead3c2 100644
--- a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-pt-rBR/strings.xml
+++ b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-pt-rBR/strings.xml
@@ -17,5 +17,5 @@
 
 <resources xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
-    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"Recorte de cascata"</string>
+    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"Design cachoeira"</string>
 </resources>
diff --git a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-pt/strings.xml b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-pt/strings.xml
index f80fcad..3ead3c2 100644
--- a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-pt/strings.xml
+++ b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-pt/strings.xml
@@ -17,5 +17,5 @@
 
 <resources xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
-    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"Recorte de cascata"</string>
+    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"Design cachoeira"</string>
 </resources>
diff --git a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-ta/strings.xml b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-ta/strings.xml
new file mode 100644
index 0000000..85d32c3
--- /dev/null
+++ b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-ta/strings.xml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!-- 
+  ~ Copyright (C) 2020 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+   -->
+
+<resources xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
+    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"வாட்டர்ஃபால் கட்அவுட்"</string>
+</resources>
diff --git a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-th/strings.xml b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-th/strings.xml
index 6c39a7f..e9b5e1f 100644
--- a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-th/strings.xml
+++ b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-th/strings.xml
@@ -17,5 +17,5 @@
 
 <resources xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
-    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"คัตเอาท์ Waterfall"</string>
+    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"ขอบจอโค้ง"</string>
 </resources>
diff --git a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-vi/strings.xml b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-vi/strings.xml
index 2fff027..a063e8f 100644
--- a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-vi/strings.xml
+++ b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-vi/strings.xml
@@ -17,5 +17,5 @@
 
 <resources xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
-    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"Vết cắt trên thác nước"</string>
+    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"Vết cắt thác nước"</string>
 </resources>
diff --git a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-zh-rTW/strings.xml b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-zh-rTW/strings.xml
index 109b61c..56144e5 100644
--- a/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-zh-rTW/strings.xml
+++ b/packages/overlays/DisplayCutoutEmulationWaterfallOverlay/res/values-zh-rTW/strings.xml
@@ -17,5 +17,5 @@
 
 <resources xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
-    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"瀑布裁剪圖片"</string>
+    <string name="display_cutout_emulation_overlay" msgid="3523556473422419323">"瀑布螢幕凹口"</string>
 </resources>
diff --git a/services/core/java/com/android/server/media/MediaSessionRecord.java b/services/core/java/com/android/server/media/MediaSessionRecord.java
index 9f47b34..67f9782 100644
--- a/services/core/java/com/android/server/media/MediaSessionRecord.java
+++ b/services/core/java/com/android/server/media/MediaSessionRecord.java
@@ -172,7 +172,7 @@
         mTag = tag;
         mSessionInfo = sessionInfo;
         mController = new ControllerStub();
-        mSessionToken = new MediaSession.Token(mController);
+        mSessionToken = new MediaSession.Token(ownerUid, mController);
         mSession = new SessionStub();
         mSessionCb = new SessionCb(cb);
         mService = service;
diff --git a/services/core/java/com/android/server/wm/DisplayPolicy.java b/services/core/java/com/android/server/wm/DisplayPolicy.java
index 497c473..4e872ac 100644
--- a/services/core/java/com/android/server/wm/DisplayPolicy.java
+++ b/services/core/java/com/android/server/wm/DisplayPolicy.java
@@ -1032,7 +1032,9 @@
 
                             // In Gesture Nav, navigation bar frame is larger than frame to
                             // calculate inset.
-                            if (mNavigationBarPosition == NAV_BAR_BOTTOM) {
+                            if (navigationBarPosition(displayFrames.mDisplayWidth,
+                                    displayFrames.mDisplayHeight,
+                                    displayFrames.mRotation) == NAV_BAR_BOTTOM) {
                                 sTmpRect.set(displayFrames.mUnrestricted);
                                 sTmpRect.intersectUnchecked(displayFrames.mDisplayCutoutSafe);
                                 inOutFrame.top = sTmpRect.bottom
@@ -1234,10 +1236,7 @@
      * most recent layout, so they are not guaranteed to be correct.
      *
      * @param attrs The LayoutParams of the window.
-     * @param taskBounds The bounds of the task this window is on or {@code null} if no task is
-     *                   associated with the window.
-     * @param displayFrames display frames.
-     * @param floatingStack Whether the window's stack is floating.
+     * @param windowToken The token of the window.
      * @param outFrame The frame of the window.
      * @param outContentInsets The areas covered by system windows, expressed as positive insets.
      * @param outStableInsets The areas covered by stable system windows irrespective of their
@@ -1246,8 +1245,7 @@
      * @return Whether to always consume the system bars.
      *         See {@link #areSystemBarsForcedShownLw(WindowState)}.
      */
-    public boolean getLayoutHintLw(LayoutParams attrs, Rect taskBounds,
-            DisplayFrames displayFrames, boolean floatingStack, Rect outFrame,
+    boolean getLayoutHint(LayoutParams attrs, WindowToken windowToken, Rect outFrame,
             Rect outContentInsets, Rect outStableInsets,
             DisplayCutout.ParcelableWrapper outDisplayCutout) {
         final int fl = PolicyControl.getWindowFlags(null, attrs);
@@ -1260,6 +1258,18 @@
                 && (fl & FLAG_LAYOUT_INSET_DECOR) != 0;
         final boolean screenDecor = (pfl & PRIVATE_FLAG_IS_SCREEN_DECOR) != 0;
 
+        final boolean isFixedRotationTransforming =
+                windowToken != null && windowToken.isFixedRotationTransforming();
+        final ActivityRecord activity = windowToken != null ? windowToken.asActivityRecord() : null;
+        final Task task = activity != null ? activity.getTask() : null;
+        final Rect taskBounds = isFixedRotationTransforming
+                // Use token (activity) bounds if it is rotated because its task is not rotated.
+                ? windowToken.getBounds()
+                : (task != null ? task.getBounds() : null);
+        final DisplayFrames displayFrames = isFixedRotationTransforming
+                ? windowToken.getFixedRotationTransformDisplayFrames()
+                : mDisplayContent.mDisplayFrames;
+
         if (layoutInScreenAndInsetDecor && !screenDecor) {
             if ((sysUiVis & SYSTEM_UI_FLAG_LAYOUT_HIDE_NAVIGATION) != 0
                     || (attrs.getFitInsetsTypes() & Type.navigationBars()) == 0) {
@@ -1268,15 +1278,10 @@
                 outFrame.set(displayFrames.mRestricted);
             }
 
-            final Rect sf;
-            if (floatingStack) {
-                sf = null;
-            } else {
-                sf = displayFrames.mStable;
-            }
-
+            final boolean isFloatingTask = task != null && task.isFloating();
+            final Rect sf = isFloatingTask ? null : displayFrames.mStable;
             final Rect cf;
-            if (floatingStack) {
+            if (isFloatingTask) {
                 cf = null;
             } else if ((sysUiVis & View.SYSTEM_UI_FLAG_LAYOUT_STABLE) != 0) {
                 if ((fl & FLAG_FULLSCREEN) != 0) {
@@ -1425,6 +1430,7 @@
      */
     void simulateLayoutDisplay(DisplayFrames displayFrames, InsetsState insetsState, int uiMode) {
         displayFrames.onBeginLayout();
+        insetsState.setDisplayFrame(displayFrames.mUnrestricted);
         final WindowFrames simulatedWindowFrames = new WindowFrames();
         if (mNavigationBar != null) {
             simulateLayoutDecorWindow(
diff --git a/services/core/java/com/android/server/wm/InsetsSourceProvider.java b/services/core/java/com/android/server/wm/InsetsSourceProvider.java
index 0d3f6b9..58aefdc 100644
--- a/services/core/java/com/android/server/wm/InsetsSourceProvider.java
+++ b/services/core/java/com/android/server/wm/InsetsSourceProvider.java
@@ -190,7 +190,10 @@
 
     /** @return A new source computed by the specified window frame in the given display frames. */
     InsetsSource createSimulatedSource(DisplayFrames displayFrames, WindowFrames windowFrames) {
-        final InsetsSource source = new InsetsSource(mSource);
+        // Don't copy visible frame because it might not be calculated in the provided display
+        // frames and it is not significant for this usage.
+        final InsetsSource source = new InsetsSource(mSource.getType());
+        source.setVisible(mSource.isVisible());
         mTmpRect.set(windowFrames.mFrame);
         if (mFrameProvider != null) {
             mFrameProvider.accept(displayFrames, mWin, mTmpRect);
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 0169a4f..673e7e0 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -1645,23 +1645,8 @@
                 prepareNoneTransitionForRelaunching(activity);
             }
 
-            final DisplayFrames displayFrames = displayContent.mDisplayFrames;
-            // TODO: Not sure if onDisplayInfoUpdated() call is needed.
-            final DisplayInfo displayInfo = displayContent.getDisplayInfo();
-            displayFrames.onDisplayInfoUpdated(displayInfo,
-                    displayContent.calculateDisplayCutoutForRotation(displayInfo.rotation));
-            final Rect taskBounds;
-            final boolean floatingStack;
-            if (activity != null && activity.getTask() != null) {
-                taskBounds = mTmpRect;
-                tokenActivity.getTask().getBounds(mTmpRect);
-                floatingStack = activity.getTask().isFloating();
-            } else {
-                taskBounds = null;
-                floatingStack = false;
-            }
-            if (displayPolicy.getLayoutHintLw(win.mAttrs, taskBounds, displayFrames, floatingStack,
-                    outFrame, outContentInsets, outStableInsets, outDisplayCutout)) {
+            if (displayPolicy.getLayoutHint(win.mAttrs, token, outFrame, outContentInsets,
+                    outStableInsets, outDisplayCutout)) {
                 res |= WindowManagerGlobal.ADD_FLAG_ALWAYS_CONSUME_SYSTEM_BARS;
             }
             outInsetsState.set(win.getInsetsState(),
@@ -8036,27 +8021,8 @@
                             + "could not be found!");
                 }
                 final WindowToken windowToken = dc.getWindowToken(attrs.token);
-                final ActivityRecord activity;
-                if (windowToken != null && windowToken.asActivityRecord() != null) {
-                    activity = windowToken.asActivityRecord();
-                } else {
-                    activity = null;
-                }
-                final Rect taskBounds;
-                final boolean floatingStack;
-                if (activity != null && activity.getTask() != null) {
-                    final Task task = activity.getTask();
-                    taskBounds = new Rect();
-                    task.getBounds(taskBounds);
-                    floatingStack = task.isFloating();
-                } else {
-                    taskBounds = null;
-                    floatingStack = false;
-                }
-                final DisplayFrames displayFrames = dc.mDisplayFrames;
-                final DisplayPolicy policy = dc.getDisplayPolicy();
-                policy.getLayoutHintLw(attrs, taskBounds, displayFrames, floatingStack,
-                        new Rect(), outContentInsets, outStableInsets, displayCutout);
+                dc.getDisplayPolicy().getLayoutHint(attrs, windowToken, mTmpRect /* outFrame */,
+                        outContentInsets, outStableInsets, displayCutout);
             }
         } finally {
             Binder.restoreCallingIdentity(origId);
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index e452c4a..805ef09 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -1448,6 +1448,7 @@
         }
     }
 
+    /** @return The display frames in use by this window. */
     DisplayFrames getDisplayFrames(DisplayFrames originalFrames) {
         final DisplayFrames diplayFrames = mToken.getFixedRotationTransformDisplayFrames();
         if (diplayFrames != null) {
@@ -3495,8 +3496,7 @@
      */
     void notifyInsetsChanged() {
         try {
-            mClient.insetsChanged(
-                    getDisplayContent().getInsetsPolicy().getInsetsForDispatch(this));
+            mClient.insetsChanged(getInsetsState());
         } catch (RemoteException e) {
             Slog.w(TAG, "Failed to deliver inset state change w=" + this, e);
         }
@@ -3506,9 +3506,8 @@
     public void notifyInsetsControlChanged() {
         final InsetsStateController stateController =
                 getDisplayContent().getInsetsStateController();
-        final InsetsPolicy policy = getDisplayContent().getInsetsPolicy();
         try {
-            mClient.insetsControlChanged(policy.getInsetsForDispatch(this),
+            mClient.insetsControlChanged(getInsetsState(),
                     stateController.getControlsForDispatch(this));
         } catch (RemoteException e) {
             Slog.w(TAG, "Failed to deliver inset state change", e);
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/NotificationComparatorTest.java b/services/tests/uiservicestests/src/com/android/server/notification/NotificationComparatorTest.java
index a328568..b9ffd65 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/NotificationComparatorTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/NotificationComparatorTest.java
@@ -22,8 +22,8 @@
 import static org.mockito.Matchers.anyInt;
 import static org.mockito.Matchers.anyString;
 import static org.mockito.Matchers.eq;
-import static org.mockito.Mockito.when;
 import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.when;
 
 import android.app.Notification;
 import android.app.NotificationChannel;
@@ -33,6 +33,7 @@
 import android.content.pm.PackageManager;
 import android.media.session.MediaSession;
 import android.os.Build;
+import android.os.Process;
 import android.os.UserHandle;
 import android.provider.Settings;
 import android.service.notification.StatusBarNotification;
@@ -140,7 +141,7 @@
 
         Notification n3 = new Notification.Builder(mContext, TEST_CHANNEL_ID)
                 .setStyle(new Notification.MediaStyle()
-                        .setMediaSession(new MediaSession.Token(null)))
+                        .setMediaSession(new MediaSession.Token(Process.myUid(), null)))
                 .build();
         mRecordDefaultMedia = new NotificationRecord(mContext, new StatusBarNotification(pkg2,
                 pkg2, 1, "media", uid2, uid2, n3, new UserHandle(userId),
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyLayoutTests.java b/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyLayoutTests.java
index 2b0ad89..203fa61 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyLayoutTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyLayoutTests.java
@@ -50,6 +50,7 @@
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.spy;
 
+import android.app.WindowConfiguration;
 import android.graphics.Insets;
 import android.graphics.PixelFormat;
 import android.graphics.Rect;
@@ -855,9 +856,8 @@
         final DisplayCutout.ParcelableWrapper outDisplayCutout =
                 new DisplayCutout.ParcelableWrapper();
 
-        mDisplayPolicy.getLayoutHintLw(mWindow.mAttrs, null, mFrames,
-                false /* floatingStack */, outFrame, outContentInsets, outStableInsets,
-                outDisplayCutout);
+        mDisplayPolicy.getLayoutHint(mWindow.mAttrs, null /* windowToken */, outFrame,
+                outContentInsets, outStableInsets, outDisplayCutout);
 
         assertThat(outFrame, is(mFrames.mUnrestricted));
         assertThat(outContentInsets, is(new Rect(0, STATUS_BAR_HEIGHT, 0, NAV_BAR_HEIGHT)));
@@ -874,6 +874,9 @@
         mDisplayPolicy.beginLayoutLw(mFrames, 0 /* UI mode */);
 
         final Rect taskBounds = new Rect(100, 100, 200, 200);
+        final Task task = mWindow.getTask();
+        // Force the bounds because the task may resolve different bounds from Task#setBounds.
+        task.getWindowConfiguration().setBounds(taskBounds);
 
         final Rect outFrame = new Rect();
         final Rect outContentInsets = new Rect();
@@ -881,9 +884,8 @@
         final DisplayCutout.ParcelableWrapper outDisplayCutout =
                 new DisplayCutout.ParcelableWrapper();
 
-        mDisplayPolicy.getLayoutHintLw(mWindow.mAttrs, taskBounds, mFrames,
-                false /* floatingStack */, outFrame, outContentInsets, outStableInsets,
-                outDisplayCutout);
+        mDisplayPolicy.getLayoutHint(mWindow.mAttrs, mWindow.mToken, outFrame,
+                outContentInsets, outStableInsets, outDisplayCutout);
 
         assertThat(outFrame, is(taskBounds));
         assertThat(outContentInsets, is(new Rect()));
@@ -904,15 +906,20 @@
         final Rect taskBounds = new Rect(100, mFrames.mContent.bottom + 1,
                 200, mFrames.mContent.bottom + 10);
 
+        final Task task = mWindow.getTask();
+        // Make the task floating.
+        task.setWindowingMode(WindowConfiguration.WINDOWING_MODE_FREEFORM);
+        // Force the bounds because the task may resolve different bounds from Task#setBounds.
+        task.getWindowConfiguration().setBounds(taskBounds);
+
         final Rect outFrame = new Rect();
         final Rect outContentInsets = new Rect();
         final Rect outStableInsets = new Rect();
         final DisplayCutout.ParcelableWrapper outDisplayCutout =
                 new DisplayCutout.ParcelableWrapper();
 
-        mDisplayPolicy.getLayoutHintLw(mWindow.mAttrs, taskBounds, mFrames,
-                true /* floatingStack */, outFrame, outContentInsets, outStableInsets,
-                outDisplayCutout);
+        mDisplayPolicy.getLayoutHint(mWindow.mAttrs, mWindow.mToken, outFrame, outContentInsets,
+                outStableInsets, outDisplayCutout);
 
         assertThat(outFrame, is(taskBounds));
         assertThat(outContentInsets, is(new Rect()));
@@ -939,6 +946,8 @@
         final InsetsState simulatedInsetsState = new InsetsState();
         final DisplayFrames simulatedDisplayFrames = createDisplayFrames();
         mDisplayPolicy.beginLayoutLw(mFrames, uiMode);
+        // Force the display bounds because it is not synced with display frames in policy test.
+        mDisplayContent.getWindowConfiguration().setBounds(mFrames.mUnrestricted);
         mDisplayContent.getInsetsStateController().onPostLayout();
         mDisplayPolicy.simulateLayoutDisplay(simulatedDisplayFrames, simulatedInsetsState, uiMode);
 
@@ -947,20 +956,18 @@
         final StringWriter simulatedFramesDump = new StringWriter();
         simulatedDisplayFrames.dump(prefix, new PrintWriter(simulatedFramesDump));
 
-        assertEquals(realFramesDump.toString(), simulatedFramesDump.toString());
+        assertEquals(new ToStringComparatorWrapper<>(realFramesDump),
+                new ToStringComparatorWrapper<>(simulatedFramesDump));
 
-        final StringWriter realInsetsDump = new StringWriter();
         final InsetsState realInsetsState = new InsetsState(
                 mDisplayContent.getInsetsStateController().getRawInsetsState());
         // Exclude comparing IME insets because currently the simulated layout only focuses on the
         // insets from status bar and navigation bar.
         realInsetsState.removeSource(InsetsState.ITYPE_IME);
         realInsetsState.removeSource(InsetsState.ITYPE_CAPTION_BAR);
-        realInsetsState.dump(prefix, new PrintWriter(realInsetsDump));
-        final StringWriter simulatedInsetsDump = new StringWriter();
-        simulatedInsetsState.dump(prefix, new PrintWriter(simulatedInsetsDump));
 
-        assertEquals(realInsetsDump.toString(), simulatedInsetsDump.toString());
+        assertEquals(new ToStringComparatorWrapper<>(realInsetsState),
+                new ToStringComparatorWrapper<>(simulatedInsetsState));
     }
 
     @Test
diff --git a/services/tests/wmtests/src/com/android/server/wm/SystemServiceTestsBase.java b/services/tests/wmtests/src/com/android/server/wm/SystemServiceTestsBase.java
index 7e31895..f6ed314 100644
--- a/services/tests/wmtests/src/com/android/server/wm/SystemServiceTestsBase.java
+++ b/services/tests/wmtests/src/com/android/server/wm/SystemServiceTestsBase.java
@@ -60,4 +60,26 @@
     <T> T awaitInWmLock(Callable<T> callable) {
         return mLockRule.waitForLocked(callable);
     }
+
+    /**
+     * Utility class to compare the output of T#toString. It is convenient to have readable output
+     * of assertion if the string content can represent the expected states.
+     */
+    static class ToStringComparatorWrapper<T> {
+        final T mObject;
+
+        ToStringComparatorWrapper(T object) {
+            mObject = object;
+        }
+
+        @Override
+        public boolean equals(Object obj) {
+            return mObject.toString().equals(obj.toString());
+        }
+
+        @Override
+        public String toString() {
+            return mObject.toString();
+        }
+    }
 }
diff --git a/tests/RollbackTest/NetworkStagedRollbackTest.xml b/tests/RollbackTest/NetworkStagedRollbackTest.xml
index a465a4f..2ab907a 100644
--- a/tests/RollbackTest/NetworkStagedRollbackTest.xml
+++ b/tests/RollbackTest/NetworkStagedRollbackTest.xml
@@ -19,6 +19,12 @@
         <option name="cleanup-apks" value="true" />
         <option name="test-file-name" value="RollbackTest.apk" />
     </target_preparer>
+    <target_preparer class="com.android.tradefed.targetprep.RunCommandTargetPreparer">
+        <option name="run-command" value="am broadcast -a 'com.google.android.gms.phenotype.FLAG_OVERRIDE' --es package &quot;com.google.android.gms.platformconfigurator&quot; --es user '\\*' --esa flags &quot;ModuleConfig__immediate_commit_packages&quot; --esa types &quot;bytes&quot; --esa values &quot;CgA=&quot; com.google.android.gms" />
+        <option name="run-command" value="am broadcast -a 'com.google.android.gms.phenotype.FLAG_OVERRIDE' --es package &quot;com.google.android.gms.platformconfigurator&quot; --es user '\\*' --esa flags &quot;ModuleConfig__versioned_immediate_commit_packages&quot; --esa types &quot;bytes&quot; --esa values &quot;Cm5vdGFwYWNrYWdlOgA=&quot; com.google.android.gms" />
+        <option name="teardown-command" value="am broadcast -a 'com.google.android.gms.phenotype.FLAG_OVERRIDE' --es action delete --es package &quot;com.google.android.gms.platformconfigurator&quot; --es user '\*' --esa flag &quot;ModuleConfig__immediate_commit_packages&quot; com.google.android.gms" />
+        <option name="teardown-command" value="am broadcast -a 'com.google.android.gms.phenotype.FLAG_OVERRIDE' --es action delete --es package &quot;com.google.android.gms.platformconfigurator&quot; --es user '\*' --esa flag &quot;ModuleConfig__versioned_immediate_commit_packages&quot; com.google.android.gms" />
+    </target_preparer>
     <test class="com.android.tradefed.testtype.HostTest" >
         <option name="class" value="com.android.tests.rollback.host.NetworkStagedRollbackTest" />
     </test>