Support excluding persistent UI

BUG: 333312675
Test: manual

Change-Id: Ibcd504ee471fcd6328bae51111bd362ebcf90eae
diff --git a/AconfigFlags.bp b/AconfigFlags.bp
index d3e80ae..44d3702 100644
--- a/AconfigFlags.bp
+++ b/AconfigFlags.bp
@@ -87,6 +87,7 @@
         "com.android.media.flags.performance-aconfig-java",
         "com.android.media.flags.projection-aconfig-java",
         "com.android.net.thread.platform.flags-aconfig-java",
+        "com.android.server.contextualsearch.flags-java",
         "com.android.server.flags.services-aconfig-java",
         "com.android.text.flags-aconfig-java",
         "com.android.window.flags.window-aconfig-java",
@@ -801,6 +802,19 @@
     defaults: ["framework-minus-apex-aconfig-java-defaults"],
 }
 
+// Contextual Search system service
+aconfig_declarations {
+    name: "com.android.server.contextualsearch.flags-aconfig",
+    package: "com.android.server.contextualsearch.flags",
+    srcs: ["services/contextualsearch/flags/flags.aconfig"],
+}
+
+java_aconfig_library {
+    name: "com.android.server.contextualsearch.flags-java",
+    aconfig_declarations: "com.android.server.contextualsearch.flags-aconfig",
+    defaults: ["framework-minus-apex-aconfig-java-defaults"],
+}
+
 // Server Services Flags
 aconfig_declarations {
     name: "com.android.server.flags.services-aconfig",
diff --git a/services/contextualsearch/flags/flags.aconfig b/services/contextualsearch/flags/flags.aconfig
new file mode 100644
index 0000000..6d2da06
--- /dev/null
+++ b/services/contextualsearch/flags/flags.aconfig
@@ -0,0 +1,10 @@
+package: "com.android.server.contextualsearch.flags"
+container: "system"
+
+flag {
+    name: "enable_exclude_persistent_ui"
+    namespace: "machine_learning"
+    description: "Excluding persistent UI from contextual search screenshot."
+    is_fixed_read_only: true
+    bug: "333312675"
+}
\ No newline at end of file
diff --git a/services/contextualsearch/java/com/android/server/contextualsearch/ContextualSearchManagerService.java b/services/contextualsearch/java/com/android/server/contextualsearch/ContextualSearchManagerService.java
index 16a9933..fdd1025 100644
--- a/services/contextualsearch/java/com/android/server/contextualsearch/ContextualSearchManagerService.java
+++ b/services/contextualsearch/java/com/android/server/contextualsearch/ContextualSearchManagerService.java
@@ -23,6 +23,11 @@
 import static android.content.Intent.FLAG_ACTIVITY_NO_USER_ACTION;
 import static android.content.pm.PackageManager.MATCH_FACTORY_ONLY;
 import static android.content.pm.PackageManager.PERMISSION_GRANTED;
+import static android.view.WindowManager.LayoutParams.TYPE_NAVIGATION_BAR;
+import static android.view.WindowManager.LayoutParams.TYPE_NAVIGATION_BAR_PANEL;
+import static android.view.WindowManager.LayoutParams.TYPE_STATUS_BAR;
+
+import static com.android.server.contextualsearch.flags.Flags.enableExcludePersistentUi;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -66,6 +71,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
+import java.util.Set;
 
 public class ContextualSearchManagerService extends SystemService {
 
@@ -192,7 +198,13 @@
         }
         final ScreenCapture.ScreenshotHardwareBuffer shb;
         if (mWmInternal != null) {
-            shb = mWmInternal.takeAssistScreenshot();
+            if (enableExcludePersistentUi()) {
+                shb = mWmInternal.takeAssistScreenshot(
+                        Set.of(TYPE_STATUS_BAR, TYPE_NAVIGATION_BAR, TYPE_NAVIGATION_BAR_PANEL));
+            } else {
+                shb = mWmInternal.takeAssistScreenshot(/* windowTypesToExclude= */ Set.of());
+            }
+
         } else {
             shb = null;
         }
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 00d42e0..89405b3 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -246,6 +246,7 @@
 import android.window.DisplayWindowPolicyController;
 import android.window.IDisplayAreaOrganizer;
 import android.window.ScreenCapture;
+import android.window.ScreenCapture.LayerCaptureArgs;
 import android.window.SystemPerformanceHinter;
 import android.window.TransitionRequestInfo;
 
@@ -5232,7 +5233,7 @@
     /**
      * Creates a LayerCaptureArgs object to represent the entire DisplayContent
      */
-    ScreenCapture.LayerCaptureArgs getLayerCaptureArgs() {
+    LayerCaptureArgs getLayerCaptureArgs(Set<Integer> windowTypesToExclude) {
         if (!mWmService.mPolicy.isScreenOn()) {
             if (DEBUG_SCREENSHOT) {
                 Slog.i(TAG_WM, "Attempted to take screenshot while display was off.");
@@ -5242,8 +5243,23 @@
 
         getBounds(mTmpRect);
         mTmpRect.offsetTo(0, 0);
-        return new ScreenCapture.LayerCaptureArgs.Builder(getSurfaceControl())
-                .setSourceCrop(mTmpRect).build();
+        LayerCaptureArgs.Builder builder = new LayerCaptureArgs.Builder(getSurfaceControl())
+                .setSourceCrop(mTmpRect);
+
+        if (!windowTypesToExclude.isEmpty()) {
+            ArrayList<SurfaceControl> surfaceControls = new ArrayList<>();
+            forAllWindows(
+                    window -> {
+                        if (windowTypesToExclude.contains(window.getWindowType())) {
+                            surfaceControls.add(window.mSurfaceControl);
+                        }
+                    }, true
+            );
+            if (!surfaceControls.isEmpty()) {
+                builder.setExcludeLayers(surfaceControls.toArray(new SurfaceControl[0]));
+            }
+        }
+        return builder.build();
     }
 
     @Override
diff --git a/services/core/java/com/android/server/wm/WindowManagerInternal.java b/services/core/java/com/android/server/wm/WindowManagerInternal.java
index 7e2ffd4..407b5d5 100644
--- a/services/core/java/com/android/server/wm/WindowManagerInternal.java
+++ b/services/core/java/com/android/server/wm/WindowManagerInternal.java
@@ -50,6 +50,7 @@
 import android.view.WindowManager.DisplayImePolicy;
 import android.view.inputmethod.ImeTracker;
 import android.window.ScreenCapture;
+import android.window.ScreenCapture.ScreenshotHardwareBuffer;
 
 import com.android.internal.policy.KeyInterceptionInfo;
 import com.android.server.input.InputManagerService;
@@ -1073,8 +1074,9 @@
     public abstract boolean moveFocusToAdjacentEmbeddedActivityIfNeeded();
 
     /**
-     * Returns an instance of {@link ScreenCapture.ScreenshotHardwareBuffer} containing the current
+     * Returns an instance of {@link ScreenshotHardwareBuffer} containing the current
      * screenshot.
      */
-    public abstract ScreenCapture.ScreenshotHardwareBuffer takeAssistScreenshot();
+    public abstract ScreenshotHardwareBuffer takeAssistScreenshot(
+            Set<Integer> windowTypesToExclude);
 }
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 1496ae0..07e534c 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -320,6 +320,7 @@
 import android.window.ITrustedPresentationListener;
 import android.window.InputTransferToken;
 import android.window.ScreenCapture;
+import android.window.ScreenCapture.ScreenshotHardwareBuffer;
 import android.window.SystemPerformanceHinter;
 import android.window.TaskSnapshot;
 import android.window.TrustedPresentationThresholds;
@@ -384,6 +385,7 @@
 import java.util.NoSuchElementException;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
@@ -4172,7 +4174,7 @@
     }
 
     @Nullable
-    private ScreenCapture.ScreenshotHardwareBuffer takeAssistScreenshot() {
+    private ScreenshotHardwareBuffer takeAssistScreenshot(Set<Integer> windowTypesToExclude) {
         if (!checkCallingPermission(READ_FRAME_BUFFER, "requestAssistScreenshot()")) {
             throw new SecurityException("Requires READ_FRAME_BUFFER permission");
         }
@@ -4187,11 +4189,11 @@
                 }
                 captureArgs = null;
             } else {
-                captureArgs = displayContent.getLayerCaptureArgs();
+                captureArgs = displayContent.getLayerCaptureArgs(windowTypesToExclude);
             }
         }
 
-        final ScreenCapture.ScreenshotHardwareBuffer screenshotBuffer;
+        final ScreenshotHardwareBuffer screenshotBuffer;
         if (captureArgs != null) {
             ScreenCapture.SynchronousScreenCaptureListener syncScreenCapture =
                     ScreenCapture.createSyncCaptureListener();
@@ -4217,7 +4219,8 @@
      */
     @Override
     public boolean requestAssistScreenshot(final IAssistDataReceiver receiver) {
-        final ScreenCapture.ScreenshotHardwareBuffer shb = takeAssistScreenshot();
+        final ScreenshotHardwareBuffer shb =
+                takeAssistScreenshot(/* windowTypesToExclude= */ Set.of());
         final Bitmap bm = shb != null ? shb.asBitmap() : null;
         FgThread.getHandler().post(() -> {
             try {
@@ -4272,7 +4275,7 @@
             mTmpRect.offsetTo(0, 0);
 
             final SurfaceControl sc = task.getSurfaceControl();
-            final ScreenCapture.ScreenshotHardwareBuffer buffer = ScreenCapture.captureLayers(
+            final ScreenshotHardwareBuffer buffer = ScreenCapture.captureLayers(
                     layerCaptureArgsBuilder.setLayer(sc).setSourceCrop(mTmpRect).build());
             if (buffer == null) {
                 Slog.w(TAG, "Could not get screenshot buffer for taskId: " + taskId);
@@ -8807,9 +8810,9 @@
         }
 
         @Override
-        public ScreenCapture.ScreenshotHardwareBuffer takeAssistScreenshot() {
+        public ScreenshotHardwareBuffer takeAssistScreenshot(Set<Integer> windowTypesToExclude) {
             // WMS.takeAssistScreenshot takes care of the locking.
-            return WindowManagerService.this.takeAssistScreenshot();
+            return WindowManagerService.this.takeAssistScreenshot(windowTypesToExclude);
         }
     }
 
diff --git a/services/voiceinteraction/java/com/android/server/voiceinteraction/VoiceInteractionManagerService.java b/services/voiceinteraction/java/com/android/server/voiceinteraction/VoiceInteractionManagerService.java
index 9729c68..e5d7b40e 100644
--- a/services/voiceinteraction/java/com/android/server/voiceinteraction/VoiceInteractionManagerService.java
+++ b/services/voiceinteraction/java/com/android/server/voiceinteraction/VoiceInteractionManagerService.java
@@ -127,6 +127,7 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Objects;
+import java.util.Set;
 import java.util.concurrent.Executor;
 
 /**
@@ -2739,7 +2740,7 @@
                 }
             }
             final ScreenCapture.ScreenshotHardwareBuffer shb =
-                    mWmInternal.takeAssistScreenshot();
+                    mWmInternal.takeAssistScreenshot(/* windowTypesToExclude= */ Set.of());
             final Bitmap bm = shb != null ? shb.asBitmap() : null;
             // Now that everything is fetched, putting it in the launchIntent.
             if (bm != null) {