Expose task id and activity component in task snapshot and assist structure

- This would allow the AiAi service to cross-reference assist data with
  streaming text data and snapshots
- Pre-fill the task id/activity when receiving the assist data from the
  activity and remove unnecessary autofill santization down the line

Bug: 117268952
Test: adb shell dumpsys window all
Test: atest CtsAutoFillServiceTestCases
Test: atest CtsAssistTestCases

Change-Id: I0d0d2c85426777cc77397716db34b520593db100
diff --git a/core/java/android/app/ActivityManager.java b/core/java/android/app/ActivityManager.java
index 7330da3..1cf042f 100644
--- a/core/java/android/app/ActivityManager.java
+++ b/core/java/android/app/ActivityManager.java
@@ -1751,6 +1751,8 @@
      */
     public static class TaskSnapshot implements Parcelable {
 
+        // Top activity in task when snapshot was taken
+        private final ComponentName mTopActivityComponent;
         private final GraphicBuffer mSnapshot;
         private final int mOrientation;
         private final Rect mContentInsets;
@@ -1765,9 +1767,11 @@
         private final int mSystemUiVisibility;
         private final boolean mIsTranslucent;
 
-        public TaskSnapshot(GraphicBuffer snapshot, int orientation, Rect contentInsets,
-                boolean reducedResolution, float scale, boolean isRealSnapshot, int windowingMode,
-                int systemUiVisibility, boolean isTranslucent) {
+        public TaskSnapshot(@NonNull ComponentName topActivityComponent, GraphicBuffer snapshot,
+                int orientation, Rect contentInsets, boolean reducedResolution, float scale,
+                boolean isRealSnapshot, int windowingMode, int systemUiVisibility,
+                boolean isTranslucent) {
+            mTopActivityComponent = topActivityComponent;
             mSnapshot = snapshot;
             mOrientation = orientation;
             mContentInsets = new Rect(contentInsets);
@@ -1780,6 +1784,7 @@
         }
 
         private TaskSnapshot(Parcel source) {
+            mTopActivityComponent = ComponentName.readFromParcel(source);
             mSnapshot = source.readParcelable(null /* classLoader */);
             mOrientation = source.readInt();
             mContentInsets = source.readParcelable(null /* classLoader */);
@@ -1792,6 +1797,13 @@
         }
 
         /**
+         * @return The top activity component for the task at the point this snapshot was taken.
+         */
+        public ComponentName getTopActivityComponent() {
+            return mTopActivityComponent;
+        }
+
+        /**
          * @return The graphic buffer representing the screenshot.
          */
         @UnsupportedAppUsage
@@ -1871,6 +1883,7 @@
 
         @Override
         public void writeToParcel(Parcel dest, int flags) {
+            ComponentName.writeToParcel(mTopActivityComponent, dest);
             dest.writeParcelable(mSnapshot, 0);
             dest.writeInt(mOrientation);
             dest.writeParcelable(mContentInsets, 0);
@@ -1886,7 +1899,9 @@
         public String toString() {
             final int width = mSnapshot != null ? mSnapshot.getWidth() : 0;
             final int height = mSnapshot != null ? mSnapshot.getHeight() : 0;
-            return "TaskSnapshot{mSnapshot=" + mSnapshot + " (" + width + "x" + height + ")"
+            return "TaskSnapshot{"
+                    + " mTopActivityComponent=" + mTopActivityComponent.flattenToShortString()
+                    + " mSnapshot=" + mSnapshot + " (" + width + "x" + height + ")"
                     + " mOrientation=" + mOrientation
                     + " mContentInsets=" + mContentInsets.toShortString()
                     + " mReducedResolution=" + mReducedResolution + " mScale=" + mScale
diff --git a/core/java/android/app/assist/AssistStructure.java b/core/java/android/app/assist/AssistStructure.java
index 43f902a..cc4d4b1a 100644
--- a/core/java/android/app/assist/AssistStructure.java
+++ b/core/java/android/app/assist/AssistStructure.java
@@ -61,37 +61,39 @@
  * <a href="/guide/topics/text/autofill">Autofill Framework</a> guides.
  */
 public class AssistStructure implements Parcelable {
-    static final String TAG = "AssistStructure";
+    private static final String TAG = "AssistStructure";
 
-    static final boolean DEBUG_PARCEL = false;
-    static final boolean DEBUG_PARCEL_CHILDREN = false;
-    static final boolean DEBUG_PARCEL_TREE = false;
+    private static final boolean DEBUG_PARCEL = false;
+    private static final boolean DEBUG_PARCEL_CHILDREN = false;
+    private static final boolean DEBUG_PARCEL_TREE = false;
 
-    static final int VALIDATE_WINDOW_TOKEN = 0x11111111;
-    static final int VALIDATE_VIEW_TOKEN = 0x22222222;
+    private static final int VALIDATE_WINDOW_TOKEN = 0x11111111;
+    private static final int VALIDATE_VIEW_TOKEN = 0x22222222;
 
-    boolean mHaveData;
+    private boolean mHaveData;
 
-    ComponentName mActivityComponent;
+    // The task id and component of the activity which this assist structure is for
+    private int mTaskId;
+    private ComponentName mActivityComponent;
     private boolean mIsHomeActivity;
     private int mFlags;
     private int mAutofillFlags;
 
-    final ArrayList<WindowNode> mWindowNodes = new ArrayList<>();
+    private final ArrayList<WindowNode> mWindowNodes = new ArrayList<>();
 
-    final ArrayList<ViewNodeBuilder> mPendingAsyncChildren = new ArrayList<>();
+    private final ArrayList<ViewNodeBuilder> mPendingAsyncChildren = new ArrayList<>();
 
-    SendChannel mSendChannel;
-    IBinder mReceiveChannel;
+    private SendChannel mSendChannel;
+    private IBinder mReceiveChannel;
 
-    Rect mTmpRect = new Rect();
+    private Rect mTmpRect = new Rect();
 
-    boolean mSanitizeOnWrite = false;
+    private boolean mSanitizeOnWrite = false;
     private long mAcquisitionStartTime;
     private long mAcquisitionEndTime;
 
-    static final int TRANSACTION_XFER = Binder.FIRST_CALL_TRANSACTION+1;
-    static final String DESCRIPTOR = "android.app.AssistStructure";
+    private static final int TRANSACTION_XFER = Binder.FIRST_CALL_TRANSACTION+1;
+    private static final String DESCRIPTOR = "android.app.AssistStructure";
 
     /** @hide */
     public void setAcquisitionStartTime(long acquisitionStartTime) {
@@ -197,7 +199,6 @@
         ParcelTransferWriter(AssistStructure as, Parcel out) {
             mSanitizeOnWrite = as.mSanitizeOnWrite;
             mWriteStructure = as.waitForReady();
-            ComponentName.writeToParcel(as.mActivityComponent, out);
             out.writeInt(as.mFlags);
             out.writeInt(as.mAutofillFlags);
             out.writeLong(as.mAcquisitionStartTime);
@@ -353,7 +354,6 @@
 
         void go() {
             fetchData();
-            mActivityComponent = ComponentName.readFromParcel(mCurParcel);
             mFlags = mCurParcel.readInt();
             mAutofillFlags = mCurParcel.readInt();
             mAcquisitionStartTime = mCurParcel.readLong();
@@ -2129,7 +2129,6 @@
     /** @hide */
     public AssistStructure(Activity activity, boolean forAutoFill, int flags) {
         mHaveData = true;
-        mActivityComponent = activity.getComponentName();
         mFlags = flags;
         ArrayList<ViewRootImpl> views = WindowManagerGlobal.getInstance().getRootViews(
                 activity.getActivityToken());
@@ -2145,12 +2144,13 @@
 
     public AssistStructure() {
         mHaveData = true;
-        mActivityComponent = null;
         mFlags = 0;
     }
 
     /** @hide */
     public AssistStructure(Parcel in) {
+        mTaskId = in.readInt();
+        mActivityComponent = ComponentName.readFromParcel(in);
         mIsHomeActivity = in.readInt() == 1;
         mReceiveChannel = in.readStrongBinder();
     }
@@ -2171,7 +2171,10 @@
             Log.i(TAG, "dump(): calling ensureData() first");
             ensureData();
         }
-        Log.i(TAG, "Activity: " + mActivityComponent.flattenToShortString());
+        Log.i(TAG, "Task id: " + mTaskId);
+        Log.i(TAG, "Activity: " + (mActivityComponent != null 
+                ? mActivityComponent.flattenToShortString()
+                : null));
         Log.i(TAG, "Sanitize on write: " + mSanitizeOnWrite);
         Log.i(TAG, "Flags: " + mFlags);
         final int N = getWindowNodeCount();
@@ -2283,23 +2286,37 @@
     }
 
     /**
-     * Return the activity this AssistStructure came from.
+     * Sets the task id is associated with the activity from which this AssistStructure was
+     * generated.
+     * @hide
      */
-    public ComponentName getActivityComponent() {
-        ensureData();
-        return mActivityComponent;
+    public void setTaskId(int taskId) {
+        mTaskId = taskId;
     }
 
     /**
-     * Called by Autofill server when app forged a different value.
-     *
+     * @return The task id for the associated activity.
+     * @hide
+     */
+    public int getTaskId() {
+        return mTaskId;
+    }
+
+    /**
+     * Sets the activity that is associated with this AssistStructure.
      * @hide
      */
     public void setActivityComponent(ComponentName componentName) {
-        ensureData();
         mActivityComponent = componentName;
     }
 
+    /**
+     * Return the activity this AssistStructure came from.
+     */
+    public ComponentName getActivityComponent() {
+        return mActivityComponent;
+    }
+
     /** @hide */
     public int getFlags() {
         return mFlags;
@@ -2393,6 +2410,8 @@
 
     @Override
     public void writeToParcel(Parcel out, int flags) {
+        out.writeInt(mTaskId);
+        ComponentName.writeToParcel(mActivityComponent, out);
         out.writeInt(mIsHomeActivity ? 1 : 0);
         if (mHaveData) {
             // This object holds its data.  We want to write a send channel that the
diff --git a/core/tests/coretests/src/android/app/assist/AssistStructureTest.java b/core/tests/coretests/src/android/app/assist/AssistStructureTest.java
index fe51a39..689e683 100644
--- a/core/tests/coretests/src/android/app/assist/AssistStructureTest.java
+++ b/core/tests/coretests/src/android/app/assist/AssistStructureTest.java
@@ -133,8 +133,6 @@
     private void assertStructureWithManySmallViews(AssistStructure structure, int expectedSize) {
         int i = 0;
         try {
-            assertPackageName(structure);
-
             assertThat(structure.getWindowNodeCount()).isEqualTo(1);
 
             ViewNode rootView = structure.getWindowNodeAt(0).getRootViewNode();
@@ -188,8 +186,6 @@
 
     private void assertStructureWithOneBigView(AssistStructure structure) {
         try {
-            assertPackageName(structure);
-
             assertThat(structure.getWindowNodeCount()).isEqualTo(1);
 
             ViewNode rootView = structure.getWindowNodeAt(0).getRootViewNode();
@@ -275,12 +271,6 @@
         assertThat(hint.charAt(BIG_VIEW_SIZE - 1)).isEqualTo(BIG_VIEW_CHAR);
     }
 
-    private void assertPackageName(AssistStructure structure) {
-        assertThat(structure.getActivityComponent()).isEqualTo(
-                new ComponentName("com.android.frameworks.coretests",
-                        "android.app.assist.EmptyLayoutActivity"));
-    }
-
     private AssistStructure cloneThroughParcel(AssistStructure structure) {
         Parcel parcel = Parcel.obtain();
 
diff --git a/proto/src/task_snapshot.proto b/proto/src/task_snapshot.proto
index 65d6256..a1bbe52 100644
--- a/proto/src/task_snapshot.proto
+++ b/proto/src/task_snapshot.proto
@@ -31,4 +31,5 @@
      int32 windowing_mode = 7;
      int32 system_ui_visibility = 8;
      bool is_translucent = 9;
+     string top_activity_component = 10;
  }
\ No newline at end of file
diff --git a/services/autofill/java/com/android/server/autofill/Session.java b/services/autofill/java/com/android/server/autofill/Session.java
index 09f915e..1ff1acd 100644
--- a/services/autofill/java/com/android/server/autofill/Session.java
+++ b/services/autofill/java/com/android/server/autofill/Session.java
@@ -281,18 +281,6 @@
                     return;
                 }
 
-                // Sanitize structure before it's sent to service.
-                final ComponentName componentNameFromApp = structure.getActivityComponent();
-                if (componentNameFromApp == null || !mComponentName.getPackageName()
-                        .equals(componentNameFromApp.getPackageName())) {
-                    Slog.w(TAG, "Activity " + mComponentName + " forged different component on "
-                            + "AssistStructure: " + componentNameFromApp);
-                    structure.setActivityComponent(mComponentName);
-                    mMetricsLogger.write(newLogMaker(MetricsEvent.AUTOFILL_FORGED_COMPONENT_ATTEMPT)
-                            .addTaggedData(MetricsEvent.FIELD_AUTOFILL_FORGED_COMPONENT_NAME,
-                                    componentNameFromApp == null ? "null"
-                                            : componentNameFromApp.flattenToShortString()));
-                }
                 // Flags used to start the session.
                 int flags = structure.getFlags();
 
diff --git a/services/core/java/com/android/server/wm/ActivityRecord.java b/services/core/java/com/android/server/wm/ActivityRecord.java
index 8223693..23f8125 100644
--- a/services/core/java/com/android/server/wm/ActivityRecord.java
+++ b/services/core/java/com/android/server/wm/ActivityRecord.java
@@ -1009,8 +1009,8 @@
         updateOverrideConfiguration();
 
         mWindowContainerController = new AppWindowContainerController(taskController, appToken,
-                this, Integer.MAX_VALUE /* add on top */, info.screenOrientation, fullscreen,
-                (info.flags & FLAG_SHOW_FOR_ALL_USERS) != 0, info.configChanges,
+                realActivity, this, Integer.MAX_VALUE /* add on top */, info.screenOrientation,
+                fullscreen, (info.flags & FLAG_SHOW_FOR_ALL_USERS) != 0, info.configChanges,
                 task.voiceSession != null, mLaunchTaskBehind, isAlwaysFocusable(),
                 appInfo.targetSdkVersion, mRotationAnimationHint,
                 ActivityTaskManagerService.getInputDispatchingTimeoutLocked(this) * 1000000L);
diff --git a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
index 3359eac8..7726905 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
@@ -2658,6 +2658,9 @@
                 pae.extras.putParcelable(Intent.EXTRA_REFERRER, referrer);
             }
             if (structure != null) {
+                // Pre-fill the task/activity component for all assist data receivers
+                structure.setTaskId(pae.activity.getTask().taskId);
+                structure.setActivityComponent(pae.activity.realActivity);
                 structure.setHomeActivity(pae.isHome);
             }
             pae.haveResult = true;
diff --git a/services/core/java/com/android/server/wm/AppWindowContainerController.java b/services/core/java/com/android/server/wm/AppWindowContainerController.java
index 584c1e4..f105832 100644
--- a/services/core/java/com/android/server/wm/AppWindowContainerController.java
+++ b/services/core/java/com/android/server/wm/AppWindowContainerController.java
@@ -41,6 +41,7 @@
 import android.app.ActivityManager.TaskSnapshot;
 import android.app.ActivityOptions;
 import android.content.Intent;
+import android.content.ComponentName;
 import android.content.res.CompatibilityInfo;
 import android.content.res.Configuration;
 import android.graphics.GraphicBuffer;
@@ -199,20 +200,21 @@
     };
 
     public AppWindowContainerController(TaskWindowContainerController taskController,
-            IApplicationToken token, AppWindowContainerListener listener, int index,
-            int requestedOrientation, boolean fullscreen, boolean showForAllUsers, int configChanges,
+            IApplicationToken token, ComponentName activityComponent,
+            AppWindowContainerListener listener, int index, int requestedOrientation,
+            boolean fullscreen, boolean showForAllUsers, int configChanges,
             boolean voiceInteraction, boolean launchTaskBehind, boolean alwaysFocusable,
             int targetSdkVersion, int rotationAnimationHint, long inputDispatchingTimeoutNanos) {
-        this(taskController, token, listener, index, requestedOrientation, fullscreen,
-                showForAllUsers,
-                configChanges, voiceInteraction, launchTaskBehind, alwaysFocusable,
-                targetSdkVersion, rotationAnimationHint, inputDispatchingTimeoutNanos,
-                WindowManagerService.getInstance());
+        this(taskController, token, activityComponent, listener, index, requestedOrientation,
+                fullscreen, showForAllUsers, configChanges, voiceInteraction, launchTaskBehind,
+                alwaysFocusable, targetSdkVersion, rotationAnimationHint,
+                inputDispatchingTimeoutNanos, WindowManagerService.getInstance());
     }
 
     public AppWindowContainerController(TaskWindowContainerController taskController,
-            IApplicationToken token, AppWindowContainerListener listener, int index,
-            int requestedOrientation, boolean fullscreen, boolean showForAllUsers, int configChanges,
+            IApplicationToken token, ComponentName activityComponent,
+            AppWindowContainerListener listener, int index, int requestedOrientation,
+            boolean fullscreen, boolean showForAllUsers, int configChanges,
             boolean voiceInteraction, boolean launchTaskBehind, boolean alwaysFocusable,
             int targetSdkVersion, int rotationAnimationHint, long inputDispatchingTimeoutNanos,
             WindowManagerService service) {
@@ -233,10 +235,10 @@
                         + " controller=" + taskController);
             }
 
-            atoken = createAppWindow(mService, token, voiceInteraction, task.getDisplayContent(),
-                    inputDispatchingTimeoutNanos, fullscreen, showForAllUsers, targetSdkVersion,
-                    requestedOrientation, rotationAnimationHint, configChanges, launchTaskBehind,
-                    alwaysFocusable, this);
+            atoken = createAppWindow(mService, token, activityComponent, voiceInteraction,
+                    task.getDisplayContent(), inputDispatchingTimeoutNanos, fullscreen,
+                    showForAllUsers, targetSdkVersion, requestedOrientation, rotationAnimationHint,
+                    configChanges, launchTaskBehind, alwaysFocusable, this);
             if (DEBUG_TOKEN_MOVEMENT || DEBUG_ADD_REMOVE) Slog.v(TAG_WM, "addAppToken: " + atoken
                     + " controller=" + taskController + " at " + index);
             task.addChild(atoken, index);
@@ -245,11 +247,12 @@
 
     @VisibleForTesting
     AppWindowToken createAppWindow(WindowManagerService service, IApplicationToken token,
-            boolean voiceInteraction, DisplayContent dc, long inputDispatchingTimeoutNanos,
-            boolean fullscreen, boolean showForAllUsers, int targetSdk, int orientation,
-            int rotationAnimationHint, int configChanges, boolean launchTaskBehind,
-            boolean alwaysFocusable, AppWindowContainerController controller) {
-        return new AppWindowToken(service, token, voiceInteraction, dc,
+            ComponentName component, boolean voiceInteraction, DisplayContent dc,
+            long inputDispatchingTimeoutNanos, boolean fullscreen, boolean showForAllUsers,
+            int targetSdk, int orientation, int rotationAnimationHint, int configChanges,
+            boolean launchTaskBehind, boolean alwaysFocusable,
+            AppWindowContainerController controller) {
+        return new AppWindowToken(service, token, component, voiceInteraction, dc,
                 inputDispatchingTimeoutNanos, fullscreen, showForAllUsers, targetSdk, orientation,
                 rotationAnimationHint, configChanges, launchTaskBehind, alwaysFocusable,
                 controller);
diff --git a/services/core/java/com/android/server/wm/AppWindowToken.java b/services/core/java/com/android/server/wm/AppWindowToken.java
index 9baafcb..d30cd19 100644
--- a/services/core/java/com/android/server/wm/AppWindowToken.java
+++ b/services/core/java/com/android/server/wm/AppWindowToken.java
@@ -81,6 +81,7 @@
 
 import android.annotation.CallSuper;
 import android.app.Activity;
+import android.content.ComponentName;
 import android.content.res.Configuration;
 import android.graphics.GraphicBuffer;
 import android.graphics.Point;
@@ -130,7 +131,7 @@
 
     // Non-null only for application tokens.
     final IApplicationToken appToken;
-
+    final ComponentName mActivityComponent;
     final boolean mVoiceInteraction;
 
     /** @see WindowContainer#fillsParent() */
@@ -272,12 +273,13 @@
     /** Whether this token needs to create mAnimationBoundsLayer for cropping animations. */
     boolean mNeedsAnimationBoundsLayer;
 
-    AppWindowToken(WindowManagerService service, IApplicationToken token, boolean voiceInteraction,
-            DisplayContent dc, long inputDispatchingTimeoutNanos, boolean fullscreen,
-            boolean showForAllUsers, int targetSdk, int orientation, int rotationAnimationHint,
-            int configChanges, boolean launchTaskBehind, boolean alwaysFocusable,
+    AppWindowToken(WindowManagerService service, IApplicationToken token,
+            ComponentName activityComponent, boolean voiceInteraction, DisplayContent dc,
+            long inputDispatchingTimeoutNanos, boolean fullscreen, boolean showForAllUsers,
+            int targetSdk, int orientation, int rotationAnimationHint, int configChanges,
+            boolean launchTaskBehind, boolean alwaysFocusable,
             AppWindowContainerController controller) {
-        this(service, token, voiceInteraction, dc, fullscreen);
+        this(service, token, activityComponent, voiceInteraction, dc, fullscreen);
         setController(controller);
         mInputDispatchingTimeoutNanos = inputDispatchingTimeoutNanos;
         mShowForAllUsers = showForAllUsers;
@@ -293,11 +295,13 @@
         hiddenRequested = true;
     }
 
-    AppWindowToken(WindowManagerService service, IApplicationToken token, boolean voiceInteraction,
-            DisplayContent dc, boolean fillsParent) {
+    AppWindowToken(WindowManagerService service, IApplicationToken token,
+            ComponentName activityComponent, boolean voiceInteraction, DisplayContent dc,
+            boolean fillsParent) {
         super(service, token != null ? token.asBinder() : null, TYPE_APPLICATION, true, dc,
                 false /* ownerCanManageAppTokens */);
         appToken = token;
+        mActivityComponent = activityComponent;
         mVoiceInteraction = voiceInteraction;
         mFillsParent = fillsParent;
         mInputApplicationHandle = new InputApplicationHandle(this);
@@ -2155,6 +2159,7 @@
         if (appToken != null) {
             pw.println(prefix + "app=true mVoiceInteraction=" + mVoiceInteraction);
         }
+        pw.println(prefix + "component=" + mActivityComponent.flattenToShortString());
         pw.print(prefix); pw.print("task="); pw.println(getTask());
         pw.print(prefix); pw.print(" mFillsParent="); pw.print(mFillsParent);
                 pw.print(" mOrientation="); pw.println(mOrientation);
diff --git a/services/core/java/com/android/server/wm/TaskRecord.java b/services/core/java/com/android/server/wm/TaskRecord.java
index d697f28..d9670d6 100644
--- a/services/core/java/com/android/server/wm/TaskRecord.java
+++ b/services/core/java/com/android/server/wm/TaskRecord.java
@@ -1965,7 +1965,7 @@
                 ? reuseActivitiesReport.base.intent.getComponent()
                 : null;
         info.topActivity = reuseActivitiesReport.top != null
-                ? reuseActivitiesReport.top.intent.getComponent()
+                ? reuseActivitiesReport.top.realActivity
                 : null;
         info.origActivity = origActivity;
         info.realActivity = realActivity;
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotController.java b/services/core/java/com/android/server/wm/TaskSnapshotController.java
index b84d20d..7ab4d08 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotController.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotController.java
@@ -290,9 +290,10 @@
             return null;
         }
         final boolean isWindowTranslucent = mainWindow.getAttrs().format != PixelFormat.OPAQUE;
-        return new TaskSnapshot(buffer, appWindowToken.getConfiguration().orientation,
-                getInsets(mainWindow), isLowRamDevice /* reduced */, scaleFraction /* scale */,
-                true /* isRealSnapshot */, task.getWindowingMode(), getSystemUiVisibility(task),
+        return new TaskSnapshot(appWindowToken.mActivityComponent, buffer,
+                appWindowToken.getConfiguration().orientation, getInsets(mainWindow),
+                isLowRamDevice /* reduced */, scaleFraction /* scale */, true /* isRealSnapshot */,
+                task.getWindowingMode(), getSystemUiVisibility(task),
                 !appWindowToken.fillsParent() || isWindowTranslucent);
     }
 
@@ -382,7 +383,7 @@
         }
         // Note, the app theme snapshot is never translucent because we enforce a non-translucent
         // color above
-        return new TaskSnapshot(hwBitmap.createGraphicBufferHandle(),
+        return new TaskSnapshot(topChild.mActivityComponent, hwBitmap.createGraphicBufferHandle(),
                 topChild.getConfiguration().orientation, mainWindow.getStableInsets(),
                 ActivityManager.isLowRamDeviceStatic() /* reduced */, 1.0f /* scale */,
                 false /* isRealSnapshot */, task.getWindowingMode(), getSystemUiVisibility(task),
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotLoader.java b/services/core/java/com/android/server/wm/TaskSnapshotLoader.java
index 1410c21..0e1570b 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotLoader.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotLoader.java
@@ -21,6 +21,7 @@
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
 
 import android.app.ActivityManager.TaskSnapshot;
+import android.content.ComponentName;
 import android.graphics.Bitmap;
 import android.graphics.Bitmap.Config;
 import android.graphics.BitmapFactory;
@@ -87,7 +88,9 @@
                         + bitmapFile.getPath());
                 return null;
             }
-            return new TaskSnapshot(buffer, proto.orientation,
+            ComponentName topActivityComponent = ComponentName.unflattenFromString(
+                    proto.topActivityComponent);
+            return new TaskSnapshot(topActivityComponent, buffer, proto.orientation,
                     new Rect(proto.insetLeft, proto.insetTop, proto.insetRight, proto.insetBottom),
                     reducedResolution, reducedResolution ? REDUCED_SCALE : 1f,
                     proto.isRealSnapshot, proto.windowingMode, proto.systemUiVisibility,
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotPersister.java b/services/core/java/com/android/server/wm/TaskSnapshotPersister.java
index 6fd1795..24b5b61 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotPersister.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotPersister.java
@@ -16,7 +16,7 @@
 
 package com.android.server.wm;
 
-import static android.graphics.Bitmap.CompressFormat.*;
+import static android.graphics.Bitmap.CompressFormat.JPEG;
 
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WITH_CLASS_NAME;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
@@ -321,6 +321,7 @@
             proto.windowingMode = mSnapshot.getWindowingMode();
             proto.systemUiVisibility = mSnapshot.getSystemUiVisibility();
             proto.isTranslucent = mSnapshot.isTranslucent();
+            proto.topActivityComponent = mSnapshot.getTopActivityComponent().flattenToString();
             final byte[] bytes = TaskSnapshotProto.toByteArray(proto);
             final File file = getProtoFile(mTaskId, mUserId);
             final AtomicFile atomicFile = new AtomicFile(file);
diff --git a/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotPersisterTestBase.java b/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotPersisterTestBase.java
index 0f9b825..946ffb60 100644
--- a/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotPersisterTestBase.java
+++ b/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotPersisterTestBase.java
@@ -24,6 +24,7 @@
 import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
 
 import android.app.ActivityManager.TaskSnapshot;
+import android.content.ComponentName;
 import android.graphics.Canvas;
 import android.graphics.Color;
 import android.graphics.GraphicBuffer;
@@ -126,9 +127,9 @@
             Canvas c = buffer.lockCanvas();
             c.drawColor(Color.RED);
             buffer.unlockCanvasAndPost(c);
-            return new TaskSnapshot(buffer, ORIENTATION_PORTRAIT, TEST_INSETS,
-                    mScale < 1f /* reducedResolution */, mScale, mIsRealSnapshot, mWindowingMode,
-                    mSystemUiVisibility, mIsTranslucent);
+            return new TaskSnapshot(new ComponentName("", ""), buffer, ORIENTATION_PORTRAIT,
+                    TEST_INSETS, mScale < 1f /* reducedResolution */, mScale, mIsRealSnapshot,
+                    mWindowingMode, mSystemUiVisibility, mIsTranslucent);
         }
     }
 }
diff --git a/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotSurfaceTest.java b/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotSurfaceTest.java
index 7c16191..a569b9e 100644
--- a/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotSurfaceTest.java
+++ b/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotSurfaceTest.java
@@ -30,6 +30,7 @@
 import static org.mockito.Mockito.when;
 
 import android.app.ActivityManager.TaskSnapshot;
+import android.content.ComponentName;
 import android.graphics.Canvas;
 import android.graphics.Color;
 import android.graphics.GraphicBuffer;
@@ -60,7 +61,7 @@
             int windowFlags, Rect taskBounds) {
         final GraphicBuffer buffer = GraphicBuffer.create(width, height, PixelFormat.RGBA_8888,
                 GraphicBuffer.USAGE_SW_READ_RARELY | GraphicBuffer.USAGE_SW_WRITE_NEVER);
-        final TaskSnapshot snapshot = new TaskSnapshot(buffer,
+        final TaskSnapshot snapshot = new TaskSnapshot(new ComponentName("", ""), buffer,
                 ORIENTATION_PORTRAIT, contentInsets, false, 1.0f, true /* isRealSnapshot */,
                 WINDOWING_MODE_FULLSCREEN, 0 /* systemUiVisibility */, false /* isTranslucent */);
         mSurface = new TaskSnapshotSurface(mWm, new Window(), new Surface(), snapshot, "Test",
diff --git a/services/tests/servicestests/src/com/android/server/wm/WindowTestUtils.java b/services/tests/servicestests/src/com/android/server/wm/WindowTestUtils.java
index 9e12f02..27cf1e8 100644
--- a/services/tests/servicestests/src/com/android/server/wm/WindowTestUtils.java
+++ b/services/tests/servicestests/src/com/android/server/wm/WindowTestUtils.java
@@ -32,6 +32,7 @@
 import static org.mockito.Mockito.when;
 
 import android.app.ActivityManager;
+import android.content.ComponentName;
 import android.content.Context;
 import android.content.res.Configuration;
 import android.graphics.Rect;
@@ -152,17 +153,19 @@
         private TestAppWindowToken(DisplayContent dc) {
             super(dc.mService, new IApplicationToken.Stub() {
                 public String getName() {return null;}
-                }, false, dc, true /* fillsParent */);
+                }, new ComponentName("", ""), false, dc, true /* fillsParent */);
         }
 
         TestAppWindowToken(WindowManagerService service, IApplicationToken token,
-                boolean voiceInteraction, DisplayContent dc, long inputDispatchingTimeoutNanos,
-                boolean fullscreen, boolean showForAllUsers, int targetSdk, int orientation,
-                int rotationAnimationHint, int configChanges, boolean launchTaskBehind,
-                boolean alwaysFocusable, AppWindowContainerController controller) {
-            super(service, token, voiceInteraction, dc, inputDispatchingTimeoutNanos, fullscreen,
-                    showForAllUsers, targetSdk, orientation, rotationAnimationHint, configChanges,
-                    launchTaskBehind, alwaysFocusable, controller);
+                ComponentName activityComponent, boolean voiceInteraction, DisplayContent dc,
+                long inputDispatchingTimeoutNanos, boolean fullscreen, boolean showForAllUsers,
+                int targetSdk, int orientation, int rotationAnimationHint, int configChanges,
+                boolean launchTaskBehind, boolean alwaysFocusable,
+                AppWindowContainerController controller) {
+            super(service, token, activityComponent, voiceInteraction, dc,
+                    inputDispatchingTimeoutNanos, fullscreen, showForAllUsers, targetSdk,
+                    orientation, rotationAnimationHint, configChanges, launchTaskBehind,
+                    alwaysFocusable, controller);
         }
 
         int getWindowsCount() {
@@ -326,22 +329,24 @@
 
         TestAppWindowContainerController(TestTaskWindowContainerController taskController,
                 IApplicationToken token) {
-            super(taskController, token, null /* listener */, 0 /* index */,
-                    SCREEN_ORIENTATION_UNSPECIFIED, true /* fullscreen */,
-                    true /* showForAllUsers */, 0 /* configChanges */, false /* voiceInteraction */,
-                    false /* launchTaskBehind */, false /* alwaysFocusable */,
-                    0 /* targetSdkVersion */, 0 /* rotationAnimationHint */,
-                    0 /* inputDispatchingTimeoutNanos */, taskController.mService);
+            super(taskController, token, new ComponentName("", "") /* activityComponent */,
+                    null /* listener */, 0 /* index */, SCREEN_ORIENTATION_UNSPECIFIED,
+                    true /* fullscreen */, true /* showForAllUsers */, 0 /* configChanges */,
+                    false /* voiceInteraction */, false /* launchTaskBehind */,
+                    false /* alwaysFocusable */, 0 /* targetSdkVersion */,
+                    0 /* rotationAnimationHint */, 0 /* inputDispatchingTimeoutNanos */,
+                    taskController.mService);
             mToken = token;
         }
 
         @Override
         AppWindowToken createAppWindow(WindowManagerService service, IApplicationToken token,
-                boolean voiceInteraction, DisplayContent dc, long inputDispatchingTimeoutNanos,
+                ComponentName activityComponent, boolean voiceInteraction, DisplayContent dc,
+                long inputDispatchingTimeoutNanos,
                 boolean fullscreen, boolean showForAllUsers, int targetSdk, int orientation,
                 int rotationAnimationHint, int configChanges, boolean launchTaskBehind,
                 boolean alwaysFocusable, AppWindowContainerController controller) {
-            return new TestAppWindowToken(service, token, voiceInteraction, dc,
+            return new TestAppWindowToken(service, token, activityComponent, voiceInteraction, dc,
                     inputDispatchingTimeoutNanos, fullscreen, showForAllUsers, targetSdk,
                     orientation,
                     rotationAnimationHint, configChanges, launchTaskBehind, alwaysFocusable,