diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/ImageActionsApi.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/ImageActionsApi.java
index fd17551..e49c466 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/ImageActionsApi.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/ImageActionsApi.java
@@ -33,6 +33,7 @@
 
 import com.android.launcher3.BuildConfig;
 import com.android.quickstep.util.ImageActionUtils;
+import com.android.systemui.shared.recents.model.Task;
 
 import java.util.function.Supplier;
 
@@ -85,11 +86,11 @@
      * @param screenshotBounds the location of where the bitmap was laid out on the screen in
      *                         screen coordinates.
      * @param visibleInsets    that are used to draw the screenshot within the bounds.
-     * @param taskId           of the task that the screenshot was taken of.
+     * @param task             of the task that the screenshot was taken of.
      */
     public void saveScreenshot(Bitmap screenshot, Rect screenshotBounds,
-            Insets visibleInsets, int taskId) {
+            Insets visibleInsets, Task.TaskKey task) {
         ImageActionUtils.saveScreenshot(mSystemUiProxy, screenshot, screenshotBounds, visibleInsets,
-                taskId);
+                task);
     }
 }
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/TaskOverlayFactory.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/TaskOverlayFactory.java
index 042c542..3cfff7e 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/TaskOverlayFactory.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/TaskOverlayFactory.java
@@ -126,7 +126,7 @@
                 public void onScreenshot() {
                     if (isAllowedByPolicy) {
                         imageApi.saveScreenshot(mThumbnailView.getThumbnail(),
-                                getTaskSnapshotBounds(), getTaskSnapshotInsets(), task.key.id);
+                                getTaskSnapshotBounds(), getTaskSnapshotInsets(), task.key);
                     } else {
                         showBlockedByPolicyMessage();
                     }
diff --git a/quickstep/src/com/android/quickstep/SystemUiProxy.java b/quickstep/src/com/android/quickstep/SystemUiProxy.java
index 2ca5dc6..e7bbbd0 100644
--- a/quickstep/src/com/android/quickstep/SystemUiProxy.java
+++ b/quickstep/src/com/android/quickstep/SystemUiProxy.java
@@ -367,8 +367,17 @@
             }
         }
     }
-  
+
+    @Override
     public void handleImageBundleAsScreenshot(Bundle screenImageBundle, Rect locationInScreen,
             Insets visibleInsets, Task.TaskKey task) {
+        if (mSystemUiProxy != null) {
+            try {
+                mSystemUiProxy.handleImageBundleAsScreenshot(screenImageBundle, locationInScreen,
+                        visibleInsets, task);
+            } catch (RemoteException e) {
+                Log.w(TAG, "Failed call handleImageBundleAsScreenshot");
+            }
+        }
     }
 }
diff --git a/quickstep/src/com/android/quickstep/util/ImageActionUtils.java b/quickstep/src/com/android/quickstep/util/ImageActionUtils.java
index f5fbf28..0b48a8b 100644
--- a/quickstep/src/com/android/quickstep/util/ImageActionUtils.java
+++ b/quickstep/src/com/android/quickstep/util/ImageActionUtils.java
@@ -39,6 +39,8 @@
 
 import com.android.launcher3.BuildConfig;
 import com.android.quickstep.SystemUiProxy;
+import com.android.systemui.shared.recents.model.Task;
+import com.android.systemui.shared.recents.utilities.BitmapUtil;
 
 import java.io.File;
 import java.io.FileOutputStream;
@@ -59,8 +61,9 @@
      */
     public static void saveScreenshot(SystemUiProxy systemUiProxy, Bitmap screenshot,
             Rect screenshotBounds,
-            Insets visibleInsets, int taskId) {
-        systemUiProxy.handleImageAsScreenshot(screenshot, screenshotBounds, visibleInsets, taskId);
+            Insets visibleInsets, Task.TaskKey task) {
+        systemUiProxy.handleImageBundleAsScreenshot(BitmapUtil.hardwareBitmapToBundle(screenshot),
+                screenshotBounds, visibleInsets, task);
     }
 
     /**