3 files changed, 44 insertions, 1 deletions
diff --git a/packages/SystemUI/Android.bp b/packages/SystemUI/Android.bp
index 2fbd9ba05817..0f2e25c7025b 100644
--- a/packages/SystemUI/Android.bp
+++ b/packages/SystemUI/Android.bp
@@ -174,6 +174,9 @@ android_app {
     kotlincflags: ["-Xjvm-default=enable"],
 
     dxflags: ["--multi-dex"],
-    required: ["privapp_whitelist_com.android.systemui"],
+    required: [
+        "privapp_whitelist_com.android.systemui",
+        "checked-wm_shell_protolog.json",
+    ],
 
 }
diff --git a/packages/SystemUI/proguard.flags b/packages/SystemUI/proguard.flags
index df66bf5a1051..6c06b0a19844 100644
--- a/packages/SystemUI/proguard.flags
+++ b/packages/SystemUI/proguard.flags
@@ -41,7 +41,12 @@
     public <init>(android.content.Context);
 }
 
+# Keep the wm shell lib
 -keep class com.android.wm.shell.*
+# Keep the protolog group methods that are called by the generated code
+-keepclassmembers class com.android.wm.shell.protolog.ShellProtoLogGroup {
+    *;
+}
 
 -keep class com.android.systemui.dagger.GlobalRootComponent { *; }
 -keep class com.android.systemui.dagger.GlobalRootComponent$SysUIComponentImpl { *; }
diff --git a/packages/SystemUI/src/com/android/systemui/wmshell/WMShell.java b/packages/SystemUI/src/com/android/systemui/wmshell/WMShell.java
index 49226000e97d..e4ff1b5dcf4f 100644
--- a/packages/SystemUI/src/com/android/systemui/wmshell/WMShell.java
+++ b/packages/SystemUI/src/com/android/systemui/wmshell/WMShell.java
@@ -29,7 +29,11 @@ import com.android.systemui.shared.system.ActivityManagerWrapper;
 import com.android.systemui.shared.system.TaskStackChangeListener;
 import com.android.systemui.stackdivider.SplitScreen;
 import com.android.wm.shell.common.DisplayImeController;
+import com.android.wm.shell.protolog.ShellProtoLogImpl;
 
+import java.io.FileDescriptor;
+import java.io.PrintWriter;
+import java.util.Arrays;
 import java.util.Optional;
 
 import javax.inject.Inject;
@@ -108,4 +112,35 @@ public final class WMShell extends SystemUI {
                     }
                 });
     }
+
+    @Override
+    public void dump(FileDescriptor fd, PrintWriter pw, String[] args) {
+        // Handle commands if provided
+        for (int i = 0; i < args.length; i++) {
+            switch (args[i]) {
+                case "enable-text-logging": {
+                    String[] groups = Arrays.copyOfRange(args, i + 1, args.length);
+                    startTextLogging(groups);
+                    pw.println("Starting logging on groups: " + Arrays.toString(groups));
+                    return;
+                }
+                case "disable-text-logging": {
+                    String[] groups = Arrays.copyOfRange(args, i + 1, args.length);
+                    stopTextLogging(groups);
+                    pw.println("Stopping logging on groups: " + Arrays.toString(groups));
+                    return;
+                }
+            }
+        }
+
+        // Dump WMShell stuff here if no commands were handled
+    }
+
+    private void startTextLogging(String... groups) {
+        ShellProtoLogImpl.getSingleInstance().startTextLogging(mContext, groups);
+    }
+
+    private void stopTextLogging(String... groups) {
+        ShellProtoLogImpl.getSingleInstance().stopTextLogging(groups);
+    }
 }