diff --git a/services/core/java/com/android/server/wm/AppWindowToken.java b/services/core/java/com/android/server/wm/AppWindowToken.java
index 5838a37..71816fe 100644
--- a/services/core/java/com/android/server/wm/AppWindowToken.java
+++ b/services/core/java/com/android/server/wm/AppWindowToken.java
@@ -42,7 +42,6 @@
 import static com.android.server.wm.WindowManagerService.H.NOTIFY_ACTIVITY_DRAWN;
 import static com.android.server.wm.WindowManagerService.H.NOTIFY_STARTING_WINDOW_DRAWN;
 import static com.android.server.wm.WindowManagerService.UPDATE_FOCUS_NORMAL;
-import static com.android.server.wm.WindowStateAnimator.STACK_CLIP_NONE;
 import static com.android.server.wm.WindowManagerService.UPDATE_FOCUS_WILL_PLACE_SURFACES;
 import static com.android.server.wm.WindowManagerService.logWithStack;
 
@@ -62,12 +61,10 @@
 import android.view.IApplicationToken;
 import android.view.View;
 import android.view.WindowManager;
-import android.view.animation.Animation;
 
 import java.io.PrintWriter;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
-import java.util.function.Function;
 
 class AppTokenList extends ArrayList<AppWindowToken> {
 }
@@ -390,10 +387,10 @@
     @Override
     void removeIfPossible() {
         mIsExiting = false;
-        removeAllWindows();
+        removeAllWindowsIfPossible();
         if (mTask != null) {
             mTask.mStack.mExitingAppTokens.remove(this);
-            mTask.removeChild(this);
+            removeImmediately();
         }
     }
 
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 79d58a3..cbb1843 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -300,7 +300,7 @@
         return token.asAppWindowToken();
     }
 
-    void setWindowToken(IBinder binder, WindowToken token) {
+    void addWindowToken(IBinder binder, WindowToken token) {
         final DisplayContent dc = mService.mRoot.getWindowTokenDisplay(token);
         if (dc != null) {
             // We currently don't support adding a window token to the display if the display
@@ -333,20 +333,33 @@
     WindowToken removeWindowToken(IBinder binder) {
         final WindowToken token = mTokenMap.remove(binder);
         if (token != null && token.asAppWindowToken() == null) {
+            token.setExiting();
+        }
+        return token;
+    }
+
+    /** Changes the display the input window token is housed on to this one. */
+    void reParentWindowToken(WindowToken token) {
+        final DisplayContent prevDc = token.getDisplayContent();
+        if (prevDc == this) {
+            return;
+        }
+        if (prevDc != null && prevDc.mTokenMap.remove(token.token) != null) {
             switch (token.windowType) {
                 case TYPE_WALLPAPER:
-                    mBelowAppWindowsContainers.removeChild(token);
+                    prevDc.mBelowAppWindowsContainers.removeChild(token);
                     break;
                 case TYPE_INPUT_METHOD:
                 case TYPE_INPUT_METHOD_DIALOG:
-                    mImeWindowsContainers.removeChild(token);
+                    prevDc.mImeWindowsContainers.removeChild(token);
                     break;
                 default:
-                    mAboveAppWindowsContainers.removeChild(token);
+                    prevDc.mAboveAppWindowsContainers.removeChild(token);
                     break;
             }
         }
-        return token;
+
+        addWindowToken(token.token, token);
     }
 
     void removeAppToken(IBinder binder) {
diff --git a/services/core/java/com/android/server/wm/Task.java b/services/core/java/com/android/server/wm/Task.java
index d33ae48..fe0160a 100644
--- a/services/core/java/com/android/server/wm/Task.java
+++ b/services/core/java/com/android/server/wm/Task.java
@@ -146,11 +146,11 @@
         if (content != null) {
             content.mDimLayerController.removeDimLayerUser(this);
         }
-        getParent().removeChild(this);
+        removeImmediately();
         mService.mTaskIdToTask.delete(mTaskId);
     }
 
-    // Change to use reparenting in WC when TaskStack is switched to use WC.
+    // Change to use re-parenting in WC when TaskStack is switched to use WC.
     void moveTaskToStack(TaskStack stack, boolean toTop) {
         if (stack == mStack) {
             return;
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 51e8a56..2b33f09 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -1654,7 +1654,7 @@
         if (DEBUG_ADD_REMOVE) Slog.v(TAG_WM, "Removing " + win + " from " + token);
         // Window will already be removed from token before this post clean-up method is called.
         if (token.isEmpty()) {
-            if (!token.explicit) {
+            if (!token.mPersistOnEmpty) {
                 token.removeImmediately();
             } else if (atoken != null) {
                 // TODO: Should this be moved into AppWindowToken.removeWindow? Might go away after
@@ -2443,8 +2443,6 @@
                     return;
                 }
 
-                token.setExiting();
-
                 mInputMonitor.updateInputWindowsLw(true /*force*/);
             }
         } finally {
@@ -8924,7 +8922,7 @@
                         return;
                     }
 
-                    token.removeAllWindows();
+                    token.removeAllWindowsIfPossible();
                 }
                 WindowManagerService.this.removeWindowToken(binder, displayId);
             }
diff --git a/services/core/java/com/android/server/wm/WindowToken.java b/services/core/java/com/android/server/wm/WindowToken.java
index a2eebc3..ebf110b 100644
--- a/services/core/java/com/android/server/wm/WindowToken.java
+++ b/services/core/java/com/android/server/wm/WindowToken.java
@@ -18,25 +18,18 @@
 
 import java.util.Comparator;
 import static android.view.WindowManager.LayoutParams.FLAG_SHOW_WALLPAPER;
-import static android.view.WindowManager.LayoutParams.PRIVATE_FLAG_KEYGUARD;
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_STARTING;
 
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_ADD_REMOVE;
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_FOCUS;
-import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_LAYERS;
-import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_WALLPAPER_LIGHT;
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_WINDOW_MOVEMENT;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WITH_CLASS_NAME;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
 import static com.android.server.wm.WindowManagerService.UPDATE_FOCUS_NORMAL;
 
-import android.os.Bundle;
 import android.os.Debug;
 import android.os.IBinder;
-import android.os.RemoteException;
 import android.util.Slog;
-import android.view.DisplayInfo;
-import android.view.animation.Animation;
 
 import java.io.PrintWriter;
 
@@ -58,8 +51,8 @@
     final int windowType;
 
     // Set if this token was explicitly added by a client, so should
-    // not be removed when all windows are removed.
-    final boolean explicit;
+    // persist (not be removed) when all windows are removed.
+    boolean mPersistOnEmpty;
 
     // For printing.
     String stringName;
@@ -104,27 +97,28 @@
         return isFirstChildWindowGreaterThanSecond(newWindow, existingWindow) ? 1 : -1;
     };
 
-    WindowToken(WindowManagerService service, IBinder _token, int type, boolean _explicit,
+    WindowToken(WindowManagerService service, IBinder _token, int type, boolean persistOnEmpty,
             DisplayContent dc) {
         mService = service;
         token = _token;
         windowType = type;
-        explicit = _explicit;
+        mPersistOnEmpty = persistOnEmpty;
         onDisplayChanged(dc);
     }
 
-    void removeAllWindows() {
+    void removeAllWindowsIfPossible() {
         for (int i = mChildren.size() - 1; i >= 0; --i) {
             final WindowState win = mChildren.get(i);
-            if (DEBUG_WINDOW_MOVEMENT) Slog.w(TAG_WM, "removeAllWindows: removing win=" + win);
+            if (DEBUG_WINDOW_MOVEMENT) Slog.w(TAG_WM,
+                    "removeAllWindowsIfPossible: removing win=" + win);
             win.removeIfPossible();
-            if (mChildren.contains(win)) {
-                removeChild(win);
-            }
         }
     }
 
     void setExiting() {
+        // This token is exiting, so allow it to be removed when it no longer contains any windows.
+        mPersistOnEmpty = false;
+
         if (hidden) {
             return;
         }
@@ -297,16 +291,8 @@
     }
 
     void onDisplayChanged(DisplayContent dc) {
-        if (mDisplayContent == dc) {
-            return;
-        }
-
-        if (mDisplayContent != null) {
-            mDisplayContent.removeWindowToken(token);
-        }
+        dc.reParentWindowToken(this);
         mDisplayContent = dc;
-        mDisplayContent.setWindowToken(token, this);
-
         super.onDisplayChanged(dc);
     }
 
diff --git a/services/tests/servicestests/src/com/android/server/wm/AppWindowTokenTests.java b/services/tests/servicestests/src/com/android/server/wm/AppWindowTokenTests.java
index 207939f..06837d3 100644
--- a/services/tests/servicestests/src/com/android/server/wm/AppWindowTokenTests.java
+++ b/services/tests/servicestests/src/com/android/server/wm/AppWindowTokenTests.java
@@ -16,17 +16,12 @@
 
 package com.android.server.wm;
 
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
-import android.content.Context;
 import android.platform.test.annotations.Presubmit;
-import android.support.test.InstrumentationRegistry;
 import android.support.test.filters.SmallTest;
 import android.support.test.runner.AndroidJUnit4;
-import android.view.IWindow;
-import android.view.WindowManager;
 
 import static android.view.WindowManager.LayoutParams.FIRST_SUB_WINDOW;
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION;
@@ -49,7 +44,7 @@
 
     @Test
     public void testAddWindow_Order() throws Exception {
-        final TestAppWindowToken token = new TestAppWindowToken();
+        final TestAppWindowToken token = new TestAppWindowToken(sDisplayContent);
 
         assertEquals(0, token.getWindowsCount());
 
@@ -59,11 +54,6 @@
         final WindowState baseWin = createWindow(null, TYPE_BASE_APPLICATION, token, "baseWin");
         final WindowState win4 = createWindow(null, TYPE_APPLICATION, token, "win4");
 
-        token.addWindow(win1);
-        token.addWindow(startingWin);
-        token.addWindow(baseWin);
-        token.addWindow(win4);
-
         // Should not contain the windows that were added above.
         assertEquals(4, token.getWindowsCount());
         assertTrue(token.hasWindow(win1));
@@ -80,43 +70,17 @@
 
     @Test
     public void testFindMainWindow() throws Exception {
-        final TestAppWindowToken token = new TestAppWindowToken();
+        final TestAppWindowToken token = new TestAppWindowToken(sDisplayContent);
 
         assertNull(token.findMainWindow());
 
         final WindowState window1 = createWindow(null, TYPE_BASE_APPLICATION, token, "window1");
         final WindowState window11 = createWindow(window1, FIRST_SUB_WINDOW, token, "window11");
         final WindowState window12 = createWindow(window1, FIRST_SUB_WINDOW, token, "window12");
-        token.addWindow(window1);
         assertEquals(window1, token.findMainWindow());
         window1.mAnimatingExit = true;
         assertEquals(window1, token.findMainWindow());
         final WindowState window2 = createWindow(null, TYPE_APPLICATION_STARTING, token, "window2");
-        token.addWindow(window2);
         assertEquals(window2, token.findMainWindow());
     }
-
-    /* Used so we can gain access to some protected members of the {@link AppWindowToken} class */
-    private class TestAppWindowToken extends AppWindowToken {
-
-        TestAppWindowToken() {
-            super(sWm, null, false, sWm.getDefaultDisplayContentLocked());
-        }
-
-        int getWindowsCount() {
-            return mChildren.size();
-        }
-
-        boolean hasWindow(WindowState w) {
-            return mChildren.contains(w);
-        }
-
-        WindowState getFirstChild() {
-            return mChildren.getFirst();
-        }
-
-        WindowState getLastChild() {
-            return mChildren.getLast();
-        }
-    }
 }
diff --git a/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java b/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java
index e301b19..3a69537 100644
--- a/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java
+++ b/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java
@@ -98,9 +98,9 @@
         Assert.assertTrue("Excepted " + first + " to be greater than " + second, first > second);
     }
 
-    WindowToken createWindowToken(DisplayContent dc, int type) {
+    private WindowToken createWindowToken(DisplayContent dc, int type) {
         if (type < FIRST_APPLICATION_WINDOW || type > LAST_APPLICATION_WINDOW) {
-            return new WindowToken(sWm, mock(IBinder.class), type, false, dc);
+            return new TestWindowToken(type, dc);
         }
 
         final int stackId = sNextStackId++;
@@ -108,7 +108,7 @@
         final TaskStack stack = sWm.mStackIdToStack.get(stackId);
         final Task task = new Task(sNextTaskId++, stack, 0, sWm, null, EMPTY, false);
         stack.addTask(task, true);
-        final AppWindowToken token = new AppWindowToken(sWm, null, false, dc);
+        final TestAppWindowToken token = new TestAppWindowToken(dc);
         task.addAppToken(0, token, 0, false);
         return token;
     }
@@ -129,4 +129,48 @@
         token.addWindow(w);
         return w;
     }
+
+    /* Used so we can gain access to some protected members of the {@link WindowToken} class */
+    class TestWindowToken extends WindowToken {
+
+        TestWindowToken(int type, DisplayContent dc) {
+            this(type, dc, false /* persistOnEmpty */);
+        }
+
+        TestWindowToken(int type, DisplayContent dc, boolean persistOnEmpty) {
+            super(sWm, mock(IBinder.class), type, persistOnEmpty, dc);
+        }
+
+        int getWindowsCount() {
+            return mChildren.size();
+        }
+
+        boolean hasWindow(WindowState w) {
+            return mChildren.contains(w);
+        }
+    }
+
+    /* Used so we can gain access to some protected members of the {@link AppWindowToken} class */
+    class TestAppWindowToken extends AppWindowToken {
+
+        TestAppWindowToken(DisplayContent dc) {
+            super(sWm, null, false, dc);
+        }
+
+        int getWindowsCount() {
+            return mChildren.size();
+        }
+
+        boolean hasWindow(WindowState w) {
+            return mChildren.contains(w);
+        }
+
+        WindowState getFirstChild() {
+            return mChildren.getFirst();
+        }
+
+        WindowState getLastChild() {
+            return mChildren.getLast();
+        }
+    }
 }
diff --git a/services/tests/servicestests/src/com/android/server/wm/WindowTokenTests.java b/services/tests/servicestests/src/com/android/server/wm/WindowTokenTests.java
index d6bfa17..0c053b9 100644
--- a/services/tests/servicestests/src/com/android/server/wm/WindowTokenTests.java
+++ b/services/tests/servicestests/src/com/android/server/wm/WindowTokenTests.java
@@ -16,24 +16,19 @@
 
 package com.android.server.wm;
 
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
-import android.content.Context;
-import android.os.IBinder;
 import android.platform.test.annotations.Presubmit;
-import android.support.test.InstrumentationRegistry;
 import android.support.test.filters.SmallTest;
 import android.support.test.runner.AndroidJUnit4;
-import android.view.IWindow;
-import android.view.WindowManager;
 
-import static android.app.AppOpsManager.OP_NONE;
 import static android.view.WindowManager.LayoutParams.FIRST_SUB_WINDOW;
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION;
+import static android.view.WindowManager.LayoutParams.TYPE_TOAST;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.mock;
@@ -51,7 +46,7 @@
 
     @Test
     public void testAddWindow() throws Exception {
-        final TestWindowToken token = new TestWindowToken();
+        final TestWindowToken token = new TestWindowToken(0, sDisplayContent);
 
         assertEquals(0, token.getWindowsCount());
 
@@ -80,15 +75,13 @@
 
     @Test
     public void testChildRemoval() throws Exception {
-        final TestWindowToken token = new TestWindowToken();
-        final DisplayContent dc = sWm.getDefaultDisplayContentLocked();
+        final DisplayContent dc = sDisplayContent;
+        final TestWindowToken token = new TestWindowToken(0, dc);
 
         assertEquals(token, dc.getWindowToken(token.token));
 
         final WindowState window1 = createWindow(null, TYPE_APPLICATION, token, "window1");
         final WindowState window2 = createWindow(null, TYPE_APPLICATION, token, "window2");
-        token.addWindow(window1);
-        token.addWindow(window2);
 
         window2.removeImmediately();
         // The token should still be mapped in the display content since it still has a child.
@@ -102,17 +95,13 @@
 
     @Test
     public void testAdjustAnimLayer() throws Exception {
-        final TestWindowToken token = new TestWindowToken();
+        final TestWindowToken token = new TestWindowToken(0, sDisplayContent);
         final WindowState window1 = createWindow(null, TYPE_APPLICATION, token, "window1");
         final WindowState window11 = createWindow(window1, FIRST_SUB_WINDOW, token, "window11");
         final WindowState window12 = createWindow(window1, FIRST_SUB_WINDOW, token, "window12");
         final WindowState window2 = createWindow(null, TYPE_APPLICATION, token, "window2");
         final WindowState window3 = createWindow(null, TYPE_APPLICATION, token, "window3");
 
-        token.addWindow(window1);
-        token.addWindow(window2);
-        token.addWindow(window3);
-
         final int adj = 50;
         final int window2StartLayer = window2.mLayer = 100;
         final int window3StartLayer = window3.mLayer = 200;
@@ -126,19 +115,39 @@
         assertEquals(window3StartLayer + adj, highestLayer);
     }
 
-    /* Used so we can gain access to some protected members of the {@link WindowToken} class */
-    private class TestWindowToken extends WindowToken {
+    /**
+     * Test that a window token isn't orphaned by the system when it is requested to be removed.
+     * Tokens should only be removed from the system when all their windows are gone.
+     */
+    @Test
+    public void testTokenRemovalProcess() throws Exception {
+        final TestWindowToken token =
+                new TestWindowToken(TYPE_TOAST, sDisplayContent, true /* persistOnEmpty */);
 
-        TestWindowToken() {
-            super(sWm, mock(IBinder.class), 0, false, sWm.getDefaultDisplayContentLocked());
-        }
+        // Verify that the token is on the display
+        assertNotNull(sDisplayContent.getWindowToken(token.token));
 
-        int getWindowsCount() {
-            return mChildren.size();
-        }
+        final WindowState window1 = createWindow(null, TYPE_TOAST, token, "window1");
+        final WindowState window2 = createWindow(null, TYPE_TOAST, token, "window2");
 
-        boolean hasWindow(WindowState w) {
-            return mChildren.contains(w);
-        }
+        sDisplayContent.removeWindowToken(token.token);
+        // Verify that the token is no longer mapped on the display
+        assertNull(sDisplayContent.getWindowToken(token.token));
+        // Verify that the token is still attached to its parent
+        assertNotNull(token.getParent());
+        // Verify that the token windows are still around.
+        assertEquals(2, token.getWindowsCount());
+
+        window1.removeImmediately();
+        // Verify that the token is still attached to its parent
+        assertNotNull(token.getParent());
+        // Verify that the other token window is still around.
+        assertEquals(1, token.getWindowsCount());
+
+        window2.removeImmediately();
+        // Verify that the token is no-longer attached to its parent
+        assertNull(token.getParent());
+        // Verify that the token windows are no longer attached to it.
+        assertEquals(0, token.getWindowsCount());
     }
 }