diff --git a/services/core/java/com/android/server/wm/AppWindowToken.java b/services/core/java/com/android/server/wm/AppWindowToken.java
index 10d1d8b..ac9859d 100644
--- a/services/core/java/com/android/server/wm/AppWindowToken.java
+++ b/services/core/java/com/android/server/wm/AppWindowToken.java
@@ -445,6 +445,7 @@
 
         mService.mOpeningApps.remove(this);
         mService.mUnknownAppVisibilityController.appRemoved(this);
+        mService.mTaskSnapshotController.onAppRemoved(this);
         waitingToShow = false;
         if (mService.mClosingApps.contains(this)) {
             delayed = true;
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotCache.java b/services/core/java/com/android/server/wm/TaskSnapshotCache.java
index 994a155..c86229b 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotCache.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotCache.java
@@ -20,6 +20,8 @@
 import android.app.ActivityManager.TaskSnapshot;
 import android.util.ArrayMap;
 
+import java.io.PrintWriter;
+
 /**
  * Caches snapshots. See {@link TaskSnapshotController}.
  * <p>
@@ -27,13 +29,65 @@
  */
 class TaskSnapshotCache {
 
-    private final ArrayMap<Task, TaskSnapshot> mCache = new ArrayMap<>();
+    private final ArrayMap<AppWindowToken, Task> mAppTaskMap = new ArrayMap<>();
+    private final ArrayMap<Task, CacheEntry> mCache = new ArrayMap<>();
 
     void putSnapshot(Task task, TaskSnapshot snapshot) {
-        mCache.put(task, snapshot);
+        final CacheEntry entry = mCache.get(task);
+        if (entry != null) {
+            mAppTaskMap.remove(entry.topApp);
+        }
+        final AppWindowToken top = task.getTopChild();
+        mAppTaskMap.put(top, task);
+        mCache.put(task, new CacheEntry(snapshot, task.getTopChild()));
     }
 
     @Nullable TaskSnapshot getSnapshot(Task task) {
-        return mCache.get(task);
+        final CacheEntry entry = mCache.get(task);
+        return entry != null ? entry.snapshot : null;
+    }
+
+    /**
+     * Cleans the cache after an app window token's process died.
+     */
+    void cleanCache(AppWindowToken wtoken) {
+        final Task task = mAppTaskMap.get(wtoken);
+        if (task != null) {
+            removeEntry(task);
+        }
+    }
+
+    private void removeEntry(Task task) {
+        final CacheEntry entry = mCache.get(task);
+        if (entry != null) {
+            mAppTaskMap.remove(entry.topApp);
+            mCache.remove(task);
+        }
+    }
+
+    void dump(PrintWriter pw, String prefix) {
+        final String doublePrefix = prefix + "  ";
+        final String triplePrefix = doublePrefix + "  ";
+        pw.println(prefix + "SnapshotCache");
+        for (int i = mCache.size() - 1; i >= 0; i--) {
+            final CacheEntry entry = mCache.valueAt(i);
+            pw.println(doublePrefix + "Entry taskId=" + mCache.keyAt(i).mTaskId);
+            pw.println(triplePrefix + "topApp=" + entry.topApp);
+            pw.println(triplePrefix + "snapshot=" + entry.snapshot);
+        }
+    }
+
+    private static final class CacheEntry {
+
+        /** The snapshot. */
+        final TaskSnapshot snapshot;
+
+        /** The app token that was on top of the task when the snapshot was taken */
+        final AppWindowToken topApp;
+
+        CacheEntry(TaskSnapshot snapshot, AppWindowToken topApp) {
+            this.snapshot = snapshot;
+            this.topApp = topApp;
+        }
     }
 }
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotController.java b/services/core/java/com/android/server/wm/TaskSnapshotController.java
index 4421d61..efa72a6 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotController.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotController.java
@@ -34,6 +34,8 @@
 
 import com.android.internal.annotations.VisibleForTesting;
 
+import java.io.PrintWriter;
+
 /**
  * When an app token becomes invisible, we take a snapshot (bitmap) of the corresponding task and
  * put it into our cache. Internally we use gralloc buffers to be able to draw them wherever we
@@ -92,7 +94,7 @@
     }
 
     private TaskSnapshot snapshotTask(Task task) {
-        final AppWindowToken top = (AppWindowToken) task.getTop();
+        final AppWindowToken top = task.getTopChild();
         if (top == null) {
             return null;
         }
@@ -131,4 +133,25 @@
     private boolean canSnapshotTask(Task task) {
         return !StackId.isHomeOrRecentsStack(task.mStack.mStackId);
     }
+
+    /**
+     * Called when an {@link AppWindowToken} has been removed.
+     */
+    void onAppRemoved(AppWindowToken wtoken) {
+        // TODO: Clean from both recents and running cache.
+        mCache.cleanCache(wtoken);
+    }
+
+    /**
+     * Called when the process of an {@link AppWindowToken} has died.
+     */
+    void onAppDied(AppWindowToken wtoken) {
+
+        // TODO: Only clean from running cache.
+        mCache.cleanCache(wtoken);
+    }
+
+    void dump(PrintWriter pw, String prefix) {
+        mCache.dump(pw, prefix);
+    }
 }
diff --git a/services/core/java/com/android/server/wm/WindowContainer.java b/services/core/java/com/android/server/wm/WindowContainer.java
index fd7ea6d..15489a0 100644
--- a/services/core/java/com/android/server/wm/WindowContainer.java
+++ b/services/core/java/com/android/server/wm/WindowContainer.java
@@ -451,9 +451,9 @@
         return false;
     }
 
-    /** Returns the top child container or this container if there are no children. */
-    WindowContainer getTop() {
-        return mChildren.isEmpty() ? this : mChildren.peekLast();
+    /** Returns the top child container. */
+    E getTopChild() {
+        return mChildren.peekLast();
     }
 
     /** Returns true if there is still a removal being deferred */
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 195d4c3..b04ff42 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -7146,6 +7146,7 @@
 
         mInputMonitor.dump(pw, "  ");
         mUnknownAppVisibilityController.dump(pw, "  ");
+        mTaskSnapshotController.dump(pw, "  ");
 
         if (dumpAll) {
             pw.print("  mSystemDecorLayer="); pw.print(mSystemDecorLayer);
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index 608f056..c3b1f19 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -2297,6 +2297,9 @@
                     final WindowState win = mService.windowForClientLocked(mSession, mClient, false);
                     Slog.i(TAG, "WIN DEATH: " + win);
                     if (win != null) {
+                        if (win.mAppToken != null && win.mAppToken.findMainWindow() == win) {
+                            mService.mTaskSnapshotController.onAppDied(win.mAppToken);
+                        }
                         win.removeIfPossible(shouldKeepVisibleDeadAppWindow());
                         if (win.mAttrs.type == TYPE_DOCK_DIVIDER) {
                             // The owner of the docked divider died :( We reset the docked stack,
diff --git a/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotCacheTest.java b/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotCacheTest.java
new file mode 100644
index 0000000..48d4770
--- /dev/null
+++ b/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotCacheTest.java
@@ -0,0 +1,64 @@
+/*
+ * Copyright (C) 2017 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
+ */
+
+package com.android.server.wm;
+
+import static android.graphics.GraphicBuffer.USAGE_HW_TEXTURE;
+import static android.graphics.GraphicBuffer.USAGE_SW_READ_NEVER;
+import static android.graphics.GraphicBuffer.USAGE_SW_WRITE_NEVER;
+import static android.graphics.GraphicBuffer.create;
+import static android.view.WindowManager.LayoutParams.FIRST_APPLICATION_WINDOW;
+import static junit.framework.Assert.assertNotNull;
+import static junit.framework.Assert.assertNull;
+
+import android.app.ActivityManager.TaskSnapshot;
+import android.content.res.Configuration;
+import android.graphics.GraphicBuffer;
+import android.graphics.PixelFormat;
+import android.graphics.Rect;
+import android.platform.test.annotations.Presubmit;
+import android.support.test.filters.SmallTest;
+import android.support.test.runner.AndroidJUnit4;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+/**
+ * Test class for {@link TaskSnapshotCache}.
+ *
+ * runtest frameworks-services -c com.android.server.wm.TaskSnapshotCacheTest
+ */
+@SmallTest
+@Presubmit
+@RunWith(AndroidJUnit4.class)
+public class TaskSnapshotCacheTest extends WindowTestsBase {
+
+    @Test
+    public void testCleanCache() throws Exception {
+        TaskSnapshotCache snapshotCache = new TaskSnapshotCache();
+        final WindowState window = createWindow(null, FIRST_APPLICATION_WINDOW, "window");
+        snapshotCache.putSnapshot(window.getTask(), createSnapshot());
+        assertNotNull(snapshotCache.getSnapshot(window.getTask()));
+        snapshotCache.cleanCache(window.mAppToken);
+        assertNull(snapshotCache.getSnapshot(window.getTask()));
+    }
+
+    private TaskSnapshot createSnapshot() {
+        GraphicBuffer buffer = create(1, 1, PixelFormat.RGBA_8888,
+                USAGE_HW_TEXTURE | USAGE_SW_WRITE_NEVER | USAGE_SW_READ_NEVER);
+        return new TaskSnapshot(buffer, Configuration.ORIENTATION_PORTRAIT, new Rect());
+    }
+}