diff --git a/core/api/test-current.txt b/core/api/test-current.txt
index 4e929a8067b0f..dd97cda1b4a2c 100644
--- a/core/api/test-current.txt
+++ b/core/api/test-current.txt
@@ -382,6 +382,7 @@ package android.app {
     method @Deprecated public boolean revokeRuntimePermission(String, String, android.os.UserHandle);
     method public void syncInputTransactions();
     method public void syncInputTransactions(boolean);
+    method @Nullable public android.graphics.Bitmap takeScreenshot(@NonNull android.view.Window);
     field @NonNull public static final java.util.Set<java.lang.String> ALL_PERMISSIONS;
   }
 
diff --git a/core/java/android/app/IUiAutomationConnection.aidl b/core/java/android/app/IUiAutomationConnection.aidl
index c30bc248344c7..623af5f907498 100644
--- a/core/java/android/app/IUiAutomationConnection.aidl
+++ b/core/java/android/app/IUiAutomationConnection.aidl
@@ -20,6 +20,7 @@ import android.accessibilityservice.IAccessibilityServiceClient;
 import android.graphics.Bitmap;
 import android.graphics.Rect;
 import android.view.InputEvent;
+import android.view.SurfaceControl;
 import android.view.WindowContentFrameStats;
 import android.view.WindowAnimationFrameStats;
 import android.os.ParcelFileDescriptor;
@@ -42,6 +43,7 @@ interface IUiAutomationConnection {
     void syncInputTransactions(boolean waitForAnimations);
     boolean setRotation(int rotation);
     Bitmap takeScreenshot(in Rect crop);
+    Bitmap takeSurfaceControlScreenshot(in SurfaceControl surfaceControl);
     boolean clearWindowContentFrameStats(int windowId);
     WindowContentFrameStats getWindowContentFrameStats(int windowId);
     void clearWindowAnimationFrameStats();
diff --git a/core/java/android/app/UiAutomation.java b/core/java/android/app/UiAutomation.java
index 7d3db5e2cc491..e0b484cca8284 100644
--- a/core/java/android/app/UiAutomation.java
+++ b/core/java/android/app/UiAutomation.java
@@ -50,6 +50,10 @@ import android.view.Display;
 import android.view.InputEvent;
 import android.view.KeyEvent;
 import android.view.Surface;
+import android.view.SurfaceControl;
+import android.view.View;
+import android.view.ViewRootImpl;
+import android.view.Window;
 import android.view.WindowAnimationFrameStats;
 import android.view.WindowContentFrameStats;
 import android.view.accessibility.AccessibilityEvent;
@@ -1012,7 +1016,7 @@ public final class UiAutomation {
                 return null;
             }
         } catch (RemoteException re) {
-            Log.e(LOG_TAG, "Error while taking screnshot!", re);
+            Log.e(LOG_TAG, "Error while taking screenshot!", re);
             return null;
         }
 
@@ -1022,6 +1026,51 @@ public final class UiAutomation {
         return screenShot;
     }
 
+    /**
+     * Used to capture a screenshot of a Window. This can return null in the following cases:
+     * 1. Window content hasn't been layed out.
+     * 2. Window doesn't have a valid SurfaceControl
+     * 3. An error occurred in SurfaceFlinger when trying to take the screenshot.
+     *
+     * @param window Window to take a screenshot of
+     *
+     * @return The screenshot bitmap on success, null otherwise.
+     *
+     * @hide
+     */
+    @TestApi
+    @Nullable
+    public Bitmap takeScreenshot(@NonNull Window window) {
+        if (window == null) {
+            return null;
+        }
+
+        View decorView = window.peekDecorView();
+        if (decorView == null) {
+            return null;
+        }
+
+        ViewRootImpl viewRoot = decorView.getViewRootImpl();
+        if (viewRoot == null) {
+            return null;
+        }
+
+        SurfaceControl sc = viewRoot.getSurfaceControl();
+        if (!sc.isValid()) {
+            return null;
+        }
+
+        // Apply a sync transaction to ensure SurfaceFlinger is flushed before capturing a
+        // screenshot.
+        new SurfaceControl.Transaction().apply(true);
+        try {
+            return mUiAutomationConnection.takeSurfaceControlScreenshot(sc);
+        } catch (RemoteException re) {
+            Log.e(LOG_TAG, "Error while taking screenshot!", re);
+            return null;
+        }
+    }
+
     /**
      * Sets whether this UiAutomation to run in a "monkey" mode. Applications can query whether
      * they are executed in a "monkey" mode, i.e. run by a test framework, and avoid doing
diff --git a/core/java/android/app/UiAutomationConnection.java b/core/java/android/app/UiAutomationConnection.java
index 90210a9a27569..e693c5ea6c935 100644
--- a/core/java/android/app/UiAutomationConnection.java
+++ b/core/java/android/app/UiAutomationConnection.java
@@ -18,6 +18,7 @@ package android.app;
 
 import android.accessibilityservice.AccessibilityServiceInfo;
 import android.accessibilityservice.IAccessibilityServiceClient;
+import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.compat.annotation.UnsupportedAppUsage;
 import android.content.Context;
@@ -206,6 +207,30 @@ public final class UiAutomationConnection extends IUiAutomationConnection.Stub {
         }
     }
 
+    @Nullable
+    @Override
+    public Bitmap takeSurfaceControlScreenshot(@NonNull SurfaceControl surfaceControl) {
+        synchronized (mLock) {
+            throwIfCalledByNotTrustedUidLocked();
+            throwIfShutdownLocked();
+            throwIfNotConnectedLocked();
+        }
+
+        SurfaceControl.ScreenshotHardwareBuffer captureBuffer;
+        final long identity = Binder.clearCallingIdentity();
+        try {
+            captureBuffer = SurfaceControl.captureLayers(
+                    new SurfaceControl.LayerCaptureArgs.Builder(surfaceControl).build());
+        } finally {
+            Binder.restoreCallingIdentity(identity);
+        }
+
+        if (captureBuffer == null) {
+            return null;
+        }
+        return captureBuffer.asBitmap();
+    }
+
     @Override
     public boolean clearWindowContentFrameStats(int windowId) throws RemoteException {
         synchronized (mLock) {