Remove getLeash from WindowContainerToken

If the WindowContainer was revoked from a registered organizer, the
client could still call getLeash to system server and control the leash
for the WindowContainer. Instead, pass the leash back to the client in
onTaskAppeared and onDisplayAreaAppeared. Once the WindowContainer is
revoked from the client, the leash will reference the old
WindowContainer SurfaceControl and will not be able to control the
WindowContainer anymore.

Test: Split screen
Test: DisplayAreaOrganizerTest
Test: WindowOrganizerTest
Bug: 154558563
Change-Id: I1f6eb987a2a3fecfef912a3009ee52989c85ff4b
diff --git a/api/test-current.txt b/api/test-current.txt
index 5df84ec..bde91e3 100644
--- a/api/test-current.txt
+++ b/api/test-current.txt
@@ -5283,7 +5283,7 @@
 
   public class DisplayAreaOrganizer extends android.window.WindowOrganizer {
     ctor public DisplayAreaOrganizer();
-    method public void onDisplayAreaAppeared(@NonNull android.window.DisplayAreaInfo);
+    method public void onDisplayAreaAppeared(@NonNull android.window.DisplayAreaInfo, @NonNull android.view.SurfaceControl);
     method public void onDisplayAreaVanished(@NonNull android.window.DisplayAreaInfo);
     method @RequiresPermission(android.Manifest.permission.MANAGE_ACTIVITY_STACKS) public void registerOrganizer(int);
     field public static final int FEATURE_DEFAULT_TASK_CONTAINER = 1; // 0x1
@@ -5303,7 +5303,7 @@
     method @Nullable @RequiresPermission(android.Manifest.permission.MANAGE_ACTIVITY_STACKS) public static android.window.WindowContainerToken getImeTarget(int);
     method @Nullable @RequiresPermission(android.Manifest.permission.MANAGE_ACTIVITY_STACKS) public static java.util.List<android.app.ActivityManager.RunningTaskInfo> getRootTasks(int, @NonNull int[]);
     method @BinderThread public void onBackPressedOnTaskRoot(@NonNull android.app.ActivityManager.RunningTaskInfo);
-    method @BinderThread public void onTaskAppeared(@NonNull android.app.ActivityManager.RunningTaskInfo);
+    method @BinderThread public void onTaskAppeared(@NonNull android.app.ActivityManager.RunningTaskInfo, @NonNull android.view.SurfaceControl);
     method @BinderThread public void onTaskInfoChanged(@NonNull android.app.ActivityManager.RunningTaskInfo);
     method @BinderThread public void onTaskVanished(@NonNull android.app.ActivityManager.RunningTaskInfo);
     method @RequiresPermission(android.Manifest.permission.MANAGE_ACTIVITY_STACKS) public final void registerOrganizer(int);
@@ -5314,7 +5314,6 @@
 
   public final class WindowContainerToken implements android.os.Parcelable {
     method public int describeContents();
-    method @Nullable public android.view.SurfaceControl getLeash();
     method public void writeToParcel(@NonNull android.os.Parcel, int);
     field @NonNull public static final android.os.Parcelable.Creator<android.window.WindowContainerToken> CREATOR;
   }
diff --git a/core/java/android/window/DisplayAreaOrganizer.java b/core/java/android/window/DisplayAreaOrganizer.java
index f3ef5a0..a2fd128 100644
--- a/core/java/android/window/DisplayAreaOrganizer.java
+++ b/core/java/android/window/DisplayAreaOrganizer.java
@@ -21,6 +21,7 @@
 import android.annotation.TestApi;
 import android.os.RemoteException;
 import android.util.Singleton;
+import android.view.SurfaceControl;
 
 /**
  * Interface for WindowManager to delegate control of display areas.
@@ -64,7 +65,8 @@
         }
     }
 
-    public void onDisplayAreaAppeared(@NonNull DisplayAreaInfo displayAreaInfo) {}
+    public void onDisplayAreaAppeared(@NonNull DisplayAreaInfo displayAreaInfo,
+            @NonNull SurfaceControl leash) {}
 
     public void onDisplayAreaVanished(@NonNull DisplayAreaInfo displayAreaInfo) {}
 
@@ -76,8 +78,9 @@
     private final IDisplayAreaOrganizer mInterface = new IDisplayAreaOrganizer.Stub() {
 
         @Override
-        public void onDisplayAreaAppeared(@NonNull DisplayAreaInfo displayAreaInfo) {
-            DisplayAreaOrganizer.this.onDisplayAreaAppeared(displayAreaInfo);
+        public void onDisplayAreaAppeared(@NonNull DisplayAreaInfo displayAreaInfo,
+                @NonNull SurfaceControl leash) {
+            DisplayAreaOrganizer.this.onDisplayAreaAppeared(displayAreaInfo, leash);
         }
 
         @Override
diff --git a/core/java/android/window/IDisplayAreaOrganizer.aidl b/core/java/android/window/IDisplayAreaOrganizer.aidl
index 39a9235..5e3e5e8 100644
--- a/core/java/android/window/IDisplayAreaOrganizer.aidl
+++ b/core/java/android/window/IDisplayAreaOrganizer.aidl
@@ -17,13 +17,14 @@
 package android.window;
 
 import android.window.DisplayAreaInfo;
+import android.view.SurfaceControl;
 
 /**
  * Interface for WindowManager to delegate control of display areas.
  * {@hide}
  */
 oneway interface IDisplayAreaOrganizer {
-    void onDisplayAreaAppeared(in DisplayAreaInfo displayAreaInfo);
+    void onDisplayAreaAppeared(in DisplayAreaInfo displayAreaInfo, in SurfaceControl leash);
     void onDisplayAreaVanished(in DisplayAreaInfo displayAreaInfo);
     void onDisplayAreaInfoChanged(in DisplayAreaInfo displayAreaInfo);
 }
diff --git a/core/java/android/window/ITaskOrganizer.aidl b/core/java/android/window/ITaskOrganizer.aidl
index b4f0162..6746525 100644
--- a/core/java/android/window/ITaskOrganizer.aidl
+++ b/core/java/android/window/ITaskOrganizer.aidl
@@ -25,7 +25,14 @@
  * {@hide}
  */
 oneway interface ITaskOrganizer {
-    void onTaskAppeared(in ActivityManager.RunningTaskInfo taskInfo);
+    /**
+     * A callback when the Task is available for the registered organizer. The client is responsible
+     * for releasing the SurfaceControl in the callback.
+     *
+     * @param taskInfo The information about the Task that's available
+     * @param leash A persistent leash for this Task.
+     */
+    void onTaskAppeared(in ActivityManager.RunningTaskInfo taskInfo, in SurfaceControl leash);
     void onTaskVanished(in ActivityManager.RunningTaskInfo taskInfo);
 
     /**
diff --git a/core/java/android/window/IWindowContainerToken.aidl b/core/java/android/window/IWindowContainerToken.aidl
index 57c7abf..81a04e9 100644
--- a/core/java/android/window/IWindowContainerToken.aidl
+++ b/core/java/android/window/IWindowContainerToken.aidl
@@ -24,9 +24,4 @@
  * @hide
  */
 interface IWindowContainerToken {
-
-    /**
-     * Gets a persistent leash for this container or {@code null}.
-     */
-    SurfaceControl getLeash();
 }
diff --git a/core/java/android/window/TaskOrganizer.java b/core/java/android/window/TaskOrganizer.java
index f661d9a..1f5e533 100644
--- a/core/java/android/window/TaskOrganizer.java
+++ b/core/java/android/window/TaskOrganizer.java
@@ -24,6 +24,7 @@
 import android.app.ActivityManager;
 import android.os.RemoteException;
 import android.util.Singleton;
+import android.view.SurfaceControl;
 
 import java.util.List;
 
@@ -59,7 +60,8 @@
     }
 
     @BinderThread
-    public void onTaskAppeared(@NonNull ActivityManager.RunningTaskInfo taskInfo) {}
+    public void onTaskAppeared(@NonNull ActivityManager.RunningTaskInfo taskInfo,
+            @NonNull SurfaceControl leash) {}
 
     @BinderThread
     public void onTaskVanished(@NonNull ActivityManager.RunningTaskInfo taskInfo) {}
@@ -155,8 +157,8 @@
     private final ITaskOrganizer mInterface = new ITaskOrganizer.Stub() {
 
         @Override
-        public void onTaskAppeared(ActivityManager.RunningTaskInfo taskInfo) {
-            TaskOrganizer.this.onTaskAppeared(taskInfo);
+        public void onTaskAppeared(ActivityManager.RunningTaskInfo taskInfo, SurfaceControl leash) {
+            TaskOrganizer.this.onTaskAppeared(taskInfo, leash);
         }
 
         @Override
diff --git a/core/java/android/window/TaskOrganizerTaskEmbedder.java b/core/java/android/window/TaskOrganizerTaskEmbedder.java
index b63741e..2fb4650 100644
--- a/core/java/android/window/TaskOrganizerTaskEmbedder.java
+++ b/core/java/android/window/TaskOrganizerTaskEmbedder.java
@@ -21,7 +21,6 @@
 
 import android.app.ActivityManager;
 import android.app.ActivityOptions;
-import android.app.TaskStackListener;
 import android.content.Context;
 import android.graphics.Rect;
 import android.util.Log;
@@ -215,14 +214,14 @@
 
     private class TaskOrganizerImpl extends TaskOrganizer {
         @Override
-        public void onTaskAppeared(ActivityManager.RunningTaskInfo taskInfo) {
+        public void onTaskAppeared(ActivityManager.RunningTaskInfo taskInfo, SurfaceControl leash) {
             if (DEBUG) {
                 log("taskAppeared: " + taskInfo.taskId);
             }
 
             mTaskInfo = taskInfo;
             mTaskToken = taskInfo.token;
-            mTaskLeash = mTaskToken.getLeash();
+            mTaskLeash = leash;
             mTransaction.reparent(mTaskLeash, mSurfaceControl)
                     .show(mTaskLeash)
                     .show(mSurfaceControl)
diff --git a/core/java/android/window/WindowContainerToken.java b/core/java/android/window/WindowContainerToken.java
index 3316d0e..c92ccae 100644
--- a/core/java/android/window/WindowContainerToken.java
+++ b/core/java/android/window/WindowContainerToken.java
@@ -17,14 +17,10 @@
 package android.window;
 
 import android.annotation.NonNull;
-import android.annotation.Nullable;
 import android.annotation.TestApi;
 import android.os.IBinder;
 import android.os.Parcel;
 import android.os.Parcelable;
-import android.os.RemoteException;
-import android.view.SurfaceControl;
-import android.window.IWindowContainerToken;
 
 /**
  * Interface for a window container to communicate with the window manager. This also acts as a
@@ -45,15 +41,6 @@
         mRealToken = IWindowContainerToken.Stub.asInterface(in.readStrongBinder());
     }
 
-    @Nullable
-    public SurfaceControl getLeash() {
-        try {
-            return mRealToken.getLeash();
-        } catch (RemoteException e) {
-            throw e.rethrowFromSystemServer();
-        }
-    }
-
     /** @hide */
     public IBinder asBinder() {
         return mRealToken.asBinder();
diff --git a/packages/SystemUI/src/com/android/systemui/pip/PipTaskOrganizer.java b/packages/SystemUI/src/com/android/systemui/pip/PipTaskOrganizer.java
index 9c8fb7c..d38c481 100644
--- a/packages/SystemUI/src/com/android/systemui/pip/PipTaskOrganizer.java
+++ b/packages/SystemUI/src/com/android/systemui/pip/PipTaskOrganizer.java
@@ -250,7 +250,7 @@
     }
 
     @Override
-    public void onTaskAppeared(ActivityManager.RunningTaskInfo info) {
+    public void onTaskAppeared(ActivityManager.RunningTaskInfo info, SurfaceControl leash) {
         Objects.requireNonNull(info, "Requires RunningTaskInfo");
         final Rect destinationBounds = mPipBoundsHandler.getDestinationBounds(
                 info.topActivity, getAspectRatioOrDefault(info.pictureInPictureParams),
@@ -259,7 +259,7 @@
         mTaskInfo = info;
         mToken = mTaskInfo.token;
         mInPip = true;
-        mLeash = mToken.getLeash();
+        mLeash = leash;
 
         final Rect currentBounds = mTaskInfo.configuration.windowConfiguration.getBounds();
         mBoundsToRestore.put(mToken.asBinder(), currentBounds);
diff --git a/packages/SystemUI/src/com/android/systemui/stackdivider/SplitScreenTaskOrganizer.java b/packages/SystemUI/src/com/android/systemui/stackdivider/SplitScreenTaskOrganizer.java
index fd6722a..717edc5 100644
--- a/packages/SystemUI/src/com/android/systemui/stackdivider/SplitScreenTaskOrganizer.java
+++ b/packages/SystemUI/src/com/android/systemui/stackdivider/SplitScreenTaskOrganizer.java
@@ -83,7 +83,7 @@
     }
 
     @Override
-    public void onTaskAppeared(RunningTaskInfo taskInfo) {
+    public void onTaskAppeared(RunningTaskInfo taskInfo, SurfaceControl leash) {
         synchronized (this) {
             if (mPrimary == null || mSecondary == null) {
                 Log.w(TAG, "Received onTaskAppeared before creating root tasks " + taskInfo);
@@ -91,9 +91,9 @@
             }
 
             if (taskInfo.token.equals(mPrimary.token)) {
-                mPrimarySurface = taskInfo.token.getLeash();
+                mPrimarySurface = leash;
             } else if (taskInfo.token.equals(mSecondary.token)) {
-                mSecondarySurface = taskInfo.token.getLeash();
+                mSecondarySurface = leash;
             }
 
             if (!mSplitScreenSupported && mPrimarySurface != null && mSecondarySurface != null) {
diff --git a/services/core/java/com/android/server/wm/DisplayAreaOrganizerController.java b/services/core/java/com/android/server/wm/DisplayAreaOrganizerController.java
index 2c8c820..b6c71bb 100644
--- a/services/core/java/com/android/server/wm/DisplayAreaOrganizerController.java
+++ b/services/core/java/com/android/server/wm/DisplayAreaOrganizerController.java
@@ -21,6 +21,7 @@
 import android.os.Binder;
 import android.os.IBinder;
 import android.os.RemoteException;
+import android.view.SurfaceControl;
 import android.window.IDisplayAreaOrganizer;
 import android.window.IDisplayAreaOrganizerController;
 
@@ -113,7 +114,8 @@
 
     void onDisplayAreaAppeared(IDisplayAreaOrganizer organizer, DisplayArea da) {
         try {
-            organizer.onDisplayAreaAppeared(da.getDisplayAreaInfo());
+            SurfaceControl outSurfaceControl = new SurfaceControl(da.getSurfaceControl());
+            organizer.onDisplayAreaAppeared(da.getDisplayAreaInfo(), outSurfaceControl);
         } catch (RemoteException e) {
             // Oh well...
         }
diff --git a/services/core/java/com/android/server/wm/TaskOrganizerController.java b/services/core/java/com/android/server/wm/TaskOrganizerController.java
index 9d229df..83105c2 100644
--- a/services/core/java/com/android/server/wm/TaskOrganizerController.java
+++ b/services/core/java/com/android/server/wm/TaskOrganizerController.java
@@ -36,6 +36,7 @@
 import android.os.RemoteException;
 import android.util.Slog;
 import android.util.SparseArray;
+import android.view.SurfaceControl;
 import android.window.ITaskOrganizer;
 import android.window.ITaskOrganizerController;
 import android.window.WindowContainerToken;
@@ -112,7 +113,8 @@
             final RunningTaskInfo taskInfo = task.getTaskInfo();
             mDeferTaskOrgCallbacksConsumer.accept(() -> {
                 try {
-                    mTaskOrganizer.onTaskAppeared(taskInfo);
+                    SurfaceControl outSurfaceControl = new SurfaceControl(task.getSurfaceControl());
+                    mTaskOrganizer.onTaskAppeared(taskInfo, outSurfaceControl);
                 } catch (RemoteException e) {
                     Slog.e(TAG, "Exception sending onTaskAppeared callback", e);
                 }
diff --git a/services/core/java/com/android/server/wm/WindowContainer.java b/services/core/java/com/android/server/wm/WindowContainer.java
index 1155bb8..3f8d7b5 100644
--- a/services/core/java/com/android/server/wm/WindowContainer.java
+++ b/services/core/java/com/android/server/wm/WindowContainer.java
@@ -2557,16 +2557,6 @@
             return (RemoteToken) binder;
         }
 
-        @Override
-        public SurfaceControl getLeash() {
-            final WindowContainer wc = getContainer();
-            if (wc == null) return null;
-            // We need to copy the SurfaceControl instead of returning the original
-            // because the Parcel FLAGS PARCELABLE_WRITE_RETURN_VALUE cause SurfaceControls
-            // to release themselves.
-            return new SurfaceControl(wc.getSurfaceControl());
-        }
-
         WindowContainerToken toWindowContainerToken() {
             if (mWindowContainerToken == null) {
                 mWindowContainerToken = new WindowContainerToken(this);
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityTestsBase.java b/services/tests/wmtests/src/com/android/server/wm/ActivityTestsBase.java
index 0700f9f..2991859 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityTestsBase.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityTestsBase.java
@@ -49,6 +49,7 @@
 import android.os.Bundle;
 import android.os.UserHandle;
 import android.service.voice.IVoiceInteractionSession;
+import android.view.SurfaceControl;
 import android.window.ITaskOrganizer;
 import android.window.WindowContainerToken;
 
@@ -552,7 +553,7 @@
             mMoveToSecondaryOnEnter = move;
         }
         @Override
-        public void onTaskAppeared(ActivityManager.RunningTaskInfo info) {
+        public void onTaskAppeared(ActivityManager.RunningTaskInfo info, SurfaceControl leash) {
         }
         @Override
         public void onTaskVanished(ActivityManager.RunningTaskInfo info) {
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayAreaOrganizerTest.java b/services/tests/wmtests/src/com/android/server/wm/DisplayAreaOrganizerTest.java
index 307b40f..6a1f50d 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayAreaOrganizerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayAreaOrganizerTest.java
@@ -27,6 +27,8 @@
 import android.os.Binder;
 import android.os.RemoteException;
 import android.platform.test.annotations.Presubmit;
+import android.view.SurfaceControl;
+import android.window.DisplayAreaInfo;
 import android.window.IDisplayAreaOrganizer;
 
 import androidx.test.filters.SmallTest;
@@ -74,7 +76,8 @@
     @Test
     public void testAppearedVanished() throws RemoteException {
         IDisplayAreaOrganizer organizer = registerMockOrganizer(FEATURE_VENDOR_FIRST);
-        verify(organizer).onDisplayAreaAppeared(any());
+        verify(organizer)
+                .onDisplayAreaAppeared(any(DisplayAreaInfo.class), any(SurfaceControl.class));
 
         unregisterMockOrganizer(organizer);
         verify(organizer).onDisplayAreaVanished(any());
@@ -83,7 +86,8 @@
     @Test
     public void testChanged() throws RemoteException {
         IDisplayAreaOrganizer organizer = registerMockOrganizer(FEATURE_VENDOR_FIRST);
-        verify(organizer).onDisplayAreaAppeared(any());
+        verify(organizer)
+                .onDisplayAreaAppeared(any(DisplayAreaInfo.class), any(SurfaceControl.class));
 
         mDisplayContent.setBounds(new Rect(0, 0, 1000, 1000));
         verify(organizer).onDisplayAreaInfoChanged(any());
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
index 7d2e880..aa68c69 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
@@ -65,6 +65,7 @@
 import android.util.ArrayMap;
 import android.util.Rational;
 import android.view.Display;
+import android.view.SurfaceControl;
 import android.window.ITaskOrganizer;
 import android.window.WindowContainerTransaction;
 
@@ -133,7 +134,8 @@
         final ITaskOrganizer organizer = registerMockOrganizer();
 
         task.setTaskOrganizer(organizer);
-        verify(organizer).onTaskAppeared(any());
+        verify(organizer).onTaskAppeared(any(RunningTaskInfo.class), any(SurfaceControl.class));
+
 
         task.removeImmediately();
         verify(organizer).onTaskVanished(any());
@@ -147,11 +149,12 @@
 
         task.setTaskOrganizer(organizer);
 
-        verify(organizer, never()).onTaskAppeared(any());
+        verify(organizer, never())
+                .onTaskAppeared(any(RunningTaskInfo.class), any(SurfaceControl.class));
         task.setHasBeenVisible(true);
         assertTrue(stack.getHasBeenVisible());
 
-        verify(organizer).onTaskAppeared(any());
+        verify(organizer).onTaskAppeared(any(RunningTaskInfo.class), any(SurfaceControl.class));
 
         task.removeImmediately();
         verify(organizer).onTaskVanished(any());
@@ -167,7 +170,8 @@
         // that even though a TaskOrganizer is set remove doesn't emit
         // a vanish callback, because we never emitted appear.
         task.setTaskOrganizer(organizer);
-        verify(organizer, never()).onTaskAppeared(any());
+        verify(organizer, never())
+                .onTaskAppeared(any(RunningTaskInfo.class), any(SurfaceControl.class));
         task.removeImmediately();
         verify(organizer, never()).onTaskVanished(any());
     }
@@ -180,10 +184,10 @@
         final ITaskOrganizer organizer2 = registerMockOrganizer(WINDOWING_MODE_PINNED);
 
         task.setTaskOrganizer(organizer);
-        verify(organizer).onTaskAppeared(any());
+        verify(organizer).onTaskAppeared(any(RunningTaskInfo.class), any(SurfaceControl.class));
         task.setTaskOrganizer(organizer2);
         verify(organizer).onTaskVanished(any());
-        verify(organizer2).onTaskAppeared(any());
+        verify(organizer2).onTaskAppeared(any(RunningTaskInfo.class), any(SurfaceControl.class));
     }
 
     @Test
@@ -194,10 +198,10 @@
         final ITaskOrganizer organizer2 = registerMockOrganizer(WINDOWING_MODE_PINNED);
 
         stack.setWindowingMode(WINDOWING_MODE_MULTI_WINDOW);
-        verify(organizer).onTaskAppeared(any());
+        verify(organizer).onTaskAppeared(any(RunningTaskInfo.class), any(SurfaceControl.class));
         stack.setWindowingMode(WINDOWING_MODE_PINNED);
         verify(organizer).onTaskVanished(any());
-        verify(organizer2).onTaskAppeared(any());
+        verify(organizer2).onTaskAppeared(any(RunningTaskInfo.class), any(SurfaceControl.class));
     }
 
     @Test
@@ -207,7 +211,8 @@
         final ITaskOrganizer organizer = registerMockOrganizer();
 
         stack.setWindowingMode(WINDOWING_MODE_MULTI_WINDOW);
-        verify(organizer, never()).onTaskAppeared(any());
+        verify(organizer, never())
+                .onTaskAppeared(any(RunningTaskInfo.class), any(SurfaceControl.class));
         assertTrue(stack.isOrganized());
 
         mWm.mAtmService.mTaskOrganizerController.unregisterTaskOrganizer(organizer);
@@ -222,7 +227,7 @@
         final ITaskOrganizer organizer = registerMockOrganizer();
 
         stack.setTaskOrganizer(organizer);
-        verify(organizer).onTaskAppeared(any());
+        verify(organizer).onTaskAppeared(any(RunningTaskInfo.class), any(SurfaceControl.class));
         assertTrue(stack.isOrganized());
 
         stack.setTaskOrganizer(null);
@@ -237,7 +242,7 @@
         final ITaskOrganizer organizer = registerMockOrganizer();
 
         stack.setWindowingMode(WINDOWING_MODE_MULTI_WINDOW);
-        verify(organizer).onTaskAppeared(any());
+        verify(organizer).onTaskAppeared(any(RunningTaskInfo.class), any(SurfaceControl.class));
         assertTrue(stack.isOrganized());
 
         mWm.mAtmService.mTaskOrganizerController.unregisterTaskOrganizer(organizer);
@@ -257,7 +262,8 @@
 
         // First organizer is registered, verify a task appears when changing windowing mode
         stack.setWindowingMode(WINDOWING_MODE_MULTI_WINDOW);
-        verify(organizer, times(1)).onTaskAppeared(any());
+        verify(organizer, times(1))
+                .onTaskAppeared(any(RunningTaskInfo.class), any(SurfaceControl.class));
         assertTrue(stack.isOrganized());
 
         // Now we replace the registration and1 verify the new organizer receives tasks
@@ -265,7 +271,8 @@
         final ITaskOrganizer organizer2 = registerMockOrganizer(WINDOWING_MODE_MULTI_WINDOW);
         stack2.setWindowingMode(WINDOWING_MODE_MULTI_WINDOW);
         // One each for task and task2
-        verify(organizer2, times(2)).onTaskAppeared(any());
+        verify(organizer2, times(2))
+                .onTaskAppeared(any(RunningTaskInfo.class), any(SurfaceControl.class));
         verify(organizer2, times(0)).onTaskVanished(any());
         // One for task
         verify(organizer).onTaskVanished(any());
@@ -274,11 +281,13 @@
         // Now we unregister the second one, the first one should automatically be reregistered
         // so we verify that it's now seeing changes.
         mWm.mAtmService.mTaskOrganizerController.unregisterTaskOrganizer(organizer2);
-        verify(organizer, times(3)).onTaskAppeared(any());
+        verify(organizer, times(3))
+                .onTaskAppeared(any(RunningTaskInfo.class), any(SurfaceControl.class));
         verify(organizer2, times(2)).onTaskVanished(any());
 
         stack3.setWindowingMode(WINDOWING_MODE_MULTI_WINDOW);
-        verify(organizer, times(4)).onTaskAppeared(any());
+        verify(organizer, times(4))
+                .onTaskAppeared(any(RunningTaskInfo.class), any(SurfaceControl.class));
         verify(organizer2, times(2)).onTaskVanished(any());
         assertTrue(stack3.isOrganized());
     }
@@ -291,7 +300,8 @@
         final Task task = createTask(stack);
         final Task task2 = createTask(stack);
         stack.setWindowingMode(WINDOWING_MODE_PINNED);
-        verify(organizer, times(1)).onTaskAppeared(any());
+        verify(organizer, times(1))
+                .onTaskAppeared(any(RunningTaskInfo.class), any(SurfaceControl.class));
 
         stack.setWindowingMode(WINDOWING_MODE_FULLSCREEN);
         verify(organizer, times(1)).onTaskVanished(any());
@@ -304,7 +314,8 @@
         stack.setWindowingMode(WINDOWING_MODE_PINNED);
 
         final ITaskOrganizer organizer = registerMockOrganizer(WINDOWING_MODE_PINNED);
-        verify(organizer, times(1)).onTaskAppeared(any());
+        verify(organizer, times(1))
+                .onTaskAppeared(any(RunningTaskInfo.class), any(SurfaceControl.class));
     }
 
     @Test
@@ -459,7 +470,7 @@
     public void testTileAddRemoveChild() {
         ITaskOrganizer listener = new ITaskOrganizer.Stub() {
             @Override
-            public void onTaskAppeared(RunningTaskInfo taskInfo) { }
+            public void onTaskAppeared(RunningTaskInfo taskInfo, SurfaceControl leash) { }
 
             @Override
             public void onTaskVanished(RunningTaskInfo container) { }
@@ -515,7 +526,7 @@
         final boolean[] called = {false};
         ITaskOrganizer listener = new ITaskOrganizer.Stub() {
             @Override
-            public void onTaskAppeared(RunningTaskInfo taskInfo) { }
+            public void onTaskAppeared(RunningTaskInfo taskInfo, SurfaceControl leash) { }
 
             @Override
             public void onTaskVanished(RunningTaskInfo container) { }
@@ -577,7 +588,7 @@
         final ArrayMap<IBinder, RunningTaskInfo> lastReportedTiles = new ArrayMap<>();
         ITaskOrganizer listener = new ITaskOrganizer.Stub() {
             @Override
-            public void onTaskAppeared(RunningTaskInfo taskInfo) { }
+            public void onTaskAppeared(RunningTaskInfo taskInfo, SurfaceControl leash) { }
 
             @Override
             public void onTaskVanished(RunningTaskInfo container) { }
@@ -804,7 +815,7 @@
         RunningTaskInfo mInfo;
 
         @Override
-        public void onTaskAppeared(RunningTaskInfo info) {
+        public void onTaskAppeared(RunningTaskInfo info, SurfaceControl leash) {
             mInfo = info;
         }
         @Override
@@ -907,12 +918,14 @@
         task.setTaskOrganizer(organizer);
         // setHasBeenVisible was already called once by the set-up code.
         task.setHasBeenVisible(true);
-        verify(organizer, times(1)).onTaskAppeared(any());
+        verify(organizer, times(1))
+                .onTaskAppeared(any(RunningTaskInfo.class), any(SurfaceControl.class));
 
         task.setTaskOrganizer(null);
         verify(organizer, times(1)).onTaskVanished(any());
         task.setTaskOrganizer(organizer);
-        verify(organizer, times(2)).onTaskAppeared(any());
+        verify(organizer, times(2))
+                .onTaskAppeared(any(RunningTaskInfo.class), any(SurfaceControl.class));
 
         task.removeImmediately();
         verify(organizer, times(2)).onTaskVanished(any());
diff --git a/tests/TaskOrganizerTest/src/com/android/test/taskembed/TaskOrganizerMultiWindowTest.java b/tests/TaskOrganizerTest/src/com/android/test/taskembed/TaskOrganizerMultiWindowTest.java
index da4afec..073ae30 100644
--- a/tests/TaskOrganizerTest/src/com/android/test/taskembed/TaskOrganizerMultiWindowTest.java
+++ b/tests/TaskOrganizerTest/src/com/android/test/taskembed/TaskOrganizerMultiWindowTest.java
@@ -38,7 +38,7 @@
 import android.window.WindowContainerTransactionCallback;
 
 public class TaskOrganizerMultiWindowTest extends Activity {
-    class SplitLayout extends LinearLayout implements View.OnTouchListener {
+    static class SplitLayout extends LinearLayout implements View.OnTouchListener {
         View mView1;
         View mView2;
         View mDividerView;
@@ -94,8 +94,8 @@
     class ResizingTaskView extends TaskView {
         final Intent mIntent;
         boolean launched = false;
-        ResizingTaskView(Context c, TaskOrganizer o, int windowingMode, Intent i) {
-            super(c, o, windowingMode);
+        ResizingTaskView(Context c, Intent i) {
+            super(c);
             mIntent = i;
         }
 
@@ -120,9 +120,9 @@
         }
     }
 
-    TaskView mTaskView1;
-    TaskView mTaskView2;
-    boolean gotFirstTask = false;
+    private TaskView mTaskView1;
+    private TaskView mTaskView2;
+    private boolean mGotFirstTask = false;
 
     class Organizer extends TaskOrganizer {
         private int receivedTransactions = 0;
@@ -141,17 +141,17 @@
         };
 
         @Override
-        public void onTaskAppeared(ActivityManager.RunningTaskInfo ti) {
-            if (!gotFirstTask) {
-                mTaskView1.reparentTask(ti.token);
-                gotFirstTask = true;
+        public void onTaskAppeared(ActivityManager.RunningTaskInfo ti, SurfaceControl leash) {
+            if (!mGotFirstTask) {
+                mTaskView1.reparentTask(ti.token, leash);
+                mGotFirstTask = true;
             } else {
-                mTaskView2.reparentTask(ti.token);
+                mTaskView2.reparentTask(ti.token, leash);
             }
         }
     }
 
-    Organizer mOrganizer = new Organizer();
+    private Organizer mOrganizer = new Organizer();
 
     @Override
     protected void onCreate(Bundle savedInstanceState) {
@@ -159,10 +159,8 @@
 
         mOrganizer.registerOrganizer(WINDOWING_MODE_MULTI_WINDOW);
 
-        mTaskView1 = new ResizingTaskView(this, mOrganizer, WINDOWING_MODE_MULTI_WINDOW,
-                makeSettingsIntent());
-        mTaskView2 = new ResizingTaskView(this, mOrganizer, WINDOWING_MODE_MULTI_WINDOW,
-                makeContactsIntent());
+        mTaskView1 = new ResizingTaskView(this, makeSettingsIntent());
+        mTaskView2 = new ResizingTaskView(this, makeContactsIntent());
         View splitView = new SplitLayout(this, mTaskView1, mTaskView2);
 
         setContentView(splitView);
@@ -178,14 +176,14 @@
         intent.setFlags(Intent.FLAG_ACTIVITY_NEW_TASK | Intent.FLAG_ACTIVITY_NO_ANIMATION);
     }
 
-    Intent makeSettingsIntent() {
+    private Intent makeSettingsIntent() {
         Intent intent = new Intent();
         intent.setAction(android.provider.Settings.ACTION_SETTINGS);
         addFlags(intent);
         return intent;
     }
 
-    Intent makeContactsIntent() {
+    private Intent makeContactsIntent() {
         Intent intent = new Intent();
         intent.setAction(Intent.ACTION_MAIN);
         intent.addCategory(Intent.CATEGORY_APP_CONTACTS);
@@ -193,14 +191,14 @@
         return intent;
     }
 
-    Bundle makeLaunchOptions(int width, int height) {
+    private Bundle makeLaunchOptions(int width, int height) {
         ActivityOptions o = ActivityOptions.makeBasic();
         o.setLaunchWindowingMode(WINDOWING_MODE_MULTI_WINDOW);
         o.setLaunchBounds(new Rect(0, 0, width, height));
         return o.toBundle();
     }
 
-    void launchOrganizedActivity(Intent i, int width, int height) {
+    private void launchOrganizedActivity(Intent i, int width, int height) {
         startActivity(i, makeLaunchOptions(width, height));
     }
 }
diff --git a/tests/TaskOrganizerTest/src/com/android/test/taskembed/TaskOrganizerPipTest.java b/tests/TaskOrganizerTest/src/com/android/test/taskembed/TaskOrganizerPipTest.java
index 0a7a6d5..8fc5c5d 100644
--- a/tests/TaskOrganizerTest/src/com/android/test/taskembed/TaskOrganizerPipTest.java
+++ b/tests/TaskOrganizerTest/src/com/android/test/taskembed/TaskOrganizerPipTest.java
@@ -23,6 +23,7 @@
 import android.content.Intent;
 import android.graphics.Rect;
 import android.os.IBinder;
+import android.view.SurfaceControl;
 import android.view.ViewGroup;
 import android.view.WindowManager;
 import android.widget.FrameLayout;
@@ -30,14 +31,14 @@
 import android.window.WindowContainerTransaction;
 
 public class TaskOrganizerPipTest extends Service {
-    static final int PIP_WIDTH  = 640;
-    static final int PIP_HEIGHT = 360;
+    private static final int PIP_WIDTH  = 640;
+    private static final int PIP_HEIGHT = 360;
 
-    TaskView mTaskView;
+    private TaskView mTaskView;
 
     class Organizer extends TaskOrganizer {
-        public void onTaskAppeared(ActivityManager.RunningTaskInfo ti) {
-            mTaskView.reparentTask(ti.token);
+        public void onTaskAppeared(ActivityManager.RunningTaskInfo ti, SurfaceControl leash) {
+            mTaskView.reparentTask(ti.token, leash);
 
             final WindowContainerTransaction wct = new WindowContainerTransaction();
             wct.scheduleFinishEnterPip(ti.token, new Rect(0, 0, PIP_WIDTH, PIP_HEIGHT));
@@ -45,7 +46,7 @@
         }
     }
 
-    Organizer mOrganizer = new Organizer();
+    private Organizer mOrganizer = new Organizer();
 
     @Override
     public IBinder onBind(Intent intent) {
@@ -66,7 +67,7 @@
         FrameLayout layout = new FrameLayout(this);
         ViewGroup.LayoutParams lp =
             new ViewGroup.LayoutParams(PIP_WIDTH, PIP_HEIGHT);
-        mTaskView = new TaskView(this, mOrganizer, WINDOWING_MODE_PINNED);
+        mTaskView = new TaskView(this);
         layout.addView(mTaskView, lp);
 
         WindowManager wm = getSystemService(WindowManager.class);
diff --git a/tests/TaskOrganizerTest/src/com/android/test/taskembed/TaskView.java b/tests/TaskOrganizerTest/src/com/android/test/taskembed/TaskView.java
index aa041f2..208018c 100644
--- a/tests/TaskOrganizerTest/src/com/android/test/taskembed/TaskView.java
+++ b/tests/TaskOrganizerTest/src/com/android/test/taskembed/TaskView.java
@@ -16,34 +16,27 @@
 
 package com.android.test.taskembed;
 
-import android.app.ActivityTaskManager;
 import android.content.Context;
-import android.window.TaskOrganizer;
-import android.window.WindowContainerToken;
 import android.view.SurfaceControl;
 import android.view.SurfaceHolder;
 import android.view.SurfaceView;
-import android.window.ITaskOrganizer;
+import android.window.WindowContainerToken;
 
 /**
  * Simple SurfaceView wrapper which registers a TaskOrganizer
  * after it's Surface is ready.
  */
 class TaskView extends SurfaceView implements SurfaceHolder.Callback {
-    final TaskOrganizer mTaskOrganizer;
-    final int mWindowingMode;
     WindowContainerToken mWc;
+    private SurfaceControl mLeash;
 
-    boolean mSurfaceCreated = false;
-    boolean mNeedsReparent;
+    private boolean mSurfaceCreated = false;
+    private boolean mNeedsReparent;
 
-    TaskView(Context c, TaskOrganizer o, int windowingMode) {
+    TaskView(Context c) {
         super(c);
         getHolder().addCallback(this);
         setZOrderOnTop(true);
-
-        mTaskOrganizer = o;
-        mWindowingMode = windowingMode;
     }
 
     @Override
@@ -63,27 +56,25 @@
     public void surfaceDestroyed(SurfaceHolder holder) {
     }
 
-    void reparentTask(WindowContainerToken wc) {
+    void reparentTask(WindowContainerToken wc, SurfaceControl leash) {
         mWc = wc;
-        if (mSurfaceCreated == false) {
+        mLeash = leash;
+        if (!mSurfaceCreated) {
             mNeedsReparent = true;
         } else {
             reparentLeash();
         }
     }
 
-    void reparentLeash() {
+    private void reparentLeash() {
         SurfaceControl.Transaction t = new SurfaceControl.Transaction();
-        SurfaceControl leash = null;
-        try {
-            leash = mWc.getLeash();
-        } catch (Exception e) {
-            // System server died.. oh well
+        if (mLeash == null) {
+            return;
         }
 
-        t.reparent(leash, getSurfaceControl())
-            .setPosition(leash, 0, 0)
-            .show(leash)
+        t.reparent(mLeash, getSurfaceControl())
+            .setPosition(mLeash, 0, 0)
+            .show(mLeash)
             .apply();
     }
 }