Merge "Fix groups in WMShellFlickerServiceTests" into main
diff --git a/core/java/android/app/admin/OWNERS b/core/java/android/app/admin/OWNERS
index 308f1d6..4f3f5d9 100644
--- a/core/java/android/app/admin/OWNERS
+++ b/core/java/android/app/admin/OWNERS
@@ -1,7 +1,6 @@
 # Bug component: 142675
 # Assign bugs to device-policy-manager-triage@google.com
 
-file:WorkDeviceExperience_OWNERS
 file:EnterprisePlatformSecurity_OWNERS
 
 yamasani@google.com #{LAST_RESORT_SUGGESTION}
\ No newline at end of file
diff --git a/core/java/com/android/internal/protolog/LegacyProtoLogImpl.java b/core/java/com/android/internal/protolog/LegacyProtoLogImpl.java
index d244874..572a599 100644
--- a/core/java/com/android/internal/protolog/LegacyProtoLogImpl.java
+++ b/core/java/com/android/internal/protolog/LegacyProtoLogImpl.java
@@ -48,6 +48,7 @@
 import java.io.IOException;
 import java.io.PrintWriter;
 import java.util.ArrayList;
+import java.util.Map;
 import java.util.TreeMap;
 import java.util.stream.Collectors;
 
@@ -65,7 +66,7 @@
     private final String mLegacyViewerConfigFilename;
     private final TraceBuffer mBuffer;
     private final LegacyProtoLogViewerConfigReader mViewerConfig;
-    private final TreeMap<String, IProtoLogGroup> mLogGroups;
+    private final Map<String, IProtoLogGroup> mLogGroups = new TreeMap<>();
     private final Runnable mCacheUpdater;
     private final int mPerChunkSize;
 
@@ -74,20 +75,19 @@
     private final Object mProtoLogEnabledLock = new Object();
 
     public LegacyProtoLogImpl(String outputFile, String viewerConfigFilename,
-            TreeMap<String, IProtoLogGroup> logGroups, Runnable cacheUpdater) {
+            Runnable cacheUpdater) {
         this(new File(outputFile), viewerConfigFilename, BUFFER_CAPACITY,
-                new LegacyProtoLogViewerConfigReader(), PER_CHUNK_SIZE, logGroups, cacheUpdater);
+                new LegacyProtoLogViewerConfigReader(), PER_CHUNK_SIZE, cacheUpdater);
     }
 
     public LegacyProtoLogImpl(File file, String viewerConfigFilename, int bufferCapacity,
             LegacyProtoLogViewerConfigReader viewerConfig, int perChunkSize,
-            TreeMap<String, IProtoLogGroup> logGroups, Runnable cacheUpdater) {
+            Runnable cacheUpdater) {
         mLogFile = file;
         mBuffer = new TraceBuffer(bufferCapacity);
         mLegacyViewerConfigFilename = viewerConfigFilename;
         mViewerConfig = viewerConfig;
         mPerChunkSize = perChunkSize;
-        mLogGroups = logGroups;
         mCacheUpdater = cacheUpdater;
     }
 
@@ -97,21 +97,26 @@
     @VisibleForTesting
     @Override
     public void log(LogLevel level, IProtoLogGroup group, long messageHash, int paramsMask,
-            @Nullable String messageString, Object[] args) {
+            @Nullable Object[] args) {
         if (group.isLogToProto()) {
             logToProto(messageHash, paramsMask, args);
         }
         if (group.isLogToLogcat()) {
-            logToLogcat(group.getTag(), level, messageHash, messageString, args);
+            logToLogcat(group.getTag(), level, messageHash, args);
         }
     }
 
+    @Override
+    public void log(LogLevel logLevel, IProtoLogGroup group, String messageString, Object... args) {
+        // This will be removed very soon so no point implementing it here.
+        throw new IllegalStateException(
+                "Not implemented. Only implemented for PerfettoProtoLogImpl.");
+    }
+
     private void logToLogcat(String tag, LogLevel level, long messageHash,
-            @Nullable String messageString, Object[] args) {
+            @Nullable Object[] args) {
         String message = null;
-        if (messageString == null) {
-            messageString = mViewerConfig.getViewerString(messageHash);
-        }
+        final String messageString = mViewerConfig.getViewerString(messageHash);
         if (messageString != null) {
             if (args != null) {
                 try {
@@ -125,8 +130,10 @@
         }
         if (message == null) {
             StringBuilder builder = new StringBuilder("UNKNOWN MESSAGE (" + messageHash + ")");
-            for (Object o : args) {
-                builder.append(" ").append(o);
+            if (args != null) {
+                for (Object o : args) {
+                    builder.append(" ").append(o);
+                }
             }
             message = builder.toString();
         }
@@ -410,5 +417,12 @@
         // so we ignore the level argument to this function.
         return group.isLogToLogcat() || (group.isLogToProto() && isProtoEnabled());
     }
+
+    @Override
+    public void registerGroups(IProtoLogGroup... protoLogGroups) {
+        for (IProtoLogGroup group : protoLogGroups) {
+            mLogGroups.put(group.name(), group);
+        }
+    }
 }
 
diff --git a/core/java/com/android/internal/protolog/LogcatOnlyProtoLogImpl.java b/core/java/com/android/internal/protolog/LogcatOnlyProtoLogImpl.java
new file mode 100644
index 0000000..ebdad6d
--- /dev/null
+++ b/core/java/com/android/internal/protolog/LogcatOnlyProtoLogImpl.java
@@ -0,0 +1,87 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.internal.protolog;
+
+import static com.android.internal.protolog.ProtoLog.REQUIRE_PROTOLOGTOOL;
+
+import android.text.TextUtils;
+import android.util.Log;
+
+import com.android.internal.protolog.common.ILogger;
+import com.android.internal.protolog.common.IProtoLog;
+import com.android.internal.protolog.common.IProtoLogGroup;
+import com.android.internal.protolog.common.LogLevel;
+
+/**
+ * Class only create and used to server temporarily for when there is source code pre-processing by
+ * the ProtoLog tool, when the tracing to Perfetto flag is off, and the static REQUIRE_PROTOLOGTOOL
+ * boolean is false. In which case we simply want to log protolog message to logcat. Note, that this
+ * means that in such cases there is no real advantage of using protolog over logcat.
+ *
+ * @deprecated Should not be used. This is just a temporary class to support a legacy behavior.
+ */
+@Deprecated
+public class LogcatOnlyProtoLogImpl implements IProtoLog {
+    @Override
+    public void log(LogLevel logLevel, IProtoLogGroup group, long messageHash, int paramsMask,
+            Object[] args) {
+        throw new RuntimeException("Not supported when using LogcatOnlyProtoLogImpl");
+    }
+
+    @Override
+    public void log(LogLevel logLevel, IProtoLogGroup group, String messageString, Object[] args) {
+        if (REQUIRE_PROTOLOGTOOL) {
+            throw new RuntimeException(
+                    "REQUIRE_PROTOLOGTOOL not set to false before the first log call.");
+        }
+
+        String formattedString = TextUtils.formatSimple(messageString, args);
+        switch (logLevel) {
+            case VERBOSE -> Log.v(group.getTag(), formattedString);
+            case INFO -> Log.i(group.getTag(), formattedString);
+            case DEBUG -> Log.d(group.getTag(), formattedString);
+            case WARN -> Log.w(group.getTag(), formattedString);
+            case ERROR -> Log.e(group.getTag(), formattedString);
+            case WTF -> Log.wtf(group.getTag(), formattedString);
+        }
+    }
+
+    @Override
+    public boolean isProtoEnabled() {
+        return false;
+    }
+
+    @Override
+    public int startLoggingToLogcat(String[] groups, ILogger logger) {
+        return 0;
+    }
+
+    @Override
+    public int stopLoggingToLogcat(String[] groups, ILogger logger) {
+        return 0;
+    }
+
+    @Override
+    public boolean isEnabled(IProtoLogGroup group, LogLevel level) {
+        return true;
+    }
+
+    @Override
+    public void registerGroups(IProtoLogGroup... protoLogGroups) {
+        // Does nothing
+    }
+}
diff --git a/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java b/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
index 42fa6ac..07be700 100644
--- a/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
+++ b/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
@@ -42,11 +42,11 @@
 import static android.internal.perfetto.protos.TracePacketOuterClass.TracePacket.SEQ_NEEDS_INCREMENTAL_STATE;
 import static android.internal.perfetto.protos.TracePacketOuterClass.TracePacket.TIMESTAMP;
 
+import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.internal.perfetto.protos.Protolog.ProtoLogViewerConfig.MessageData;
 import android.os.ShellCommand;
 import android.os.SystemClock;
-import android.os.Trace;
 import android.text.TextUtils;
 import android.tracing.perfetto.DataSourceParams;
 import android.tracing.perfetto.InitArguments;
@@ -72,37 +72,45 @@
 import java.io.PrintWriter;
 import java.io.StringWriter;
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.TreeMap;
+import java.util.UUID;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 
 /**
  * A service for the ProtoLog logging system.
  */
 public class PerfettoProtoLogImpl implements IProtoLog {
     private static final String LOG_TAG = "ProtoLog";
+    public static final String NULL_STRING = "null";
     private final AtomicInteger mTracingInstances = new AtomicInteger();
 
     private final ProtoLogDataSource mDataSource = new ProtoLogDataSource(
             this::onTracingInstanceStart,
-            this::dumpTransitionTraceConfig,
+            this::onTracingFlush,
             this::onTracingInstanceStop
     );
     private final ProtoLogViewerConfigReader mViewerConfigReader;
     private final ViewerConfigInputStreamProvider mViewerConfigInputStreamProvider;
-    private final TreeMap<String, IProtoLogGroup> mLogGroups;
+    private final TreeMap<String, IProtoLogGroup> mLogGroups = new TreeMap<>();
     private final Runnable mCacheUpdater;
 
-    private final Map<LogLevel, Integer> mDefaultLogLevelCounts = new ArrayMap<>();
-    private final Map<IProtoLogGroup, Map<LogLevel, Integer>> mLogLevelCounts = new ArrayMap<>();
+    private final int[] mDefaultLogLevelCounts = new int[LogLevel.values().length];
+    private final Map<IProtoLogGroup, int[]> mLogLevelCounts = new ArrayMap<>();
+    private final Map<IProtoLogGroup, Integer> mCollectStackTraceGroupCounts = new ArrayMap<>();
 
-    private final ExecutorService mBackgroundLoggingService = Executors.newSingleThreadExecutor();
+    private final Lock mBackgroundServiceLock = new ReentrantLock();
+    private ExecutorService mBackgroundLoggingService = Executors.newSingleThreadExecutor();
 
-    public PerfettoProtoLogImpl(String viewerConfigFilePath,
-            TreeMap<String, IProtoLogGroup> logGroups, Runnable cacheUpdater) {
+    public PerfettoProtoLogImpl(String viewerConfigFilePath, Runnable cacheUpdater) {
         this(() -> {
             try {
                 return new ProtoInputStream(new FileInputStream(viewerConfigFilePath));
@@ -110,16 +118,19 @@
                 Slog.w(LOG_TAG, "Failed to load viewer config file " + viewerConfigFilePath, e);
                 return null;
             }
-        }, logGroups, cacheUpdater);
+        }, cacheUpdater);
+    }
+
+    public PerfettoProtoLogImpl() {
+        this(null, null, () -> {});
     }
 
     public PerfettoProtoLogImpl(
             ViewerConfigInputStreamProvider viewerConfigInputStreamProvider,
-            TreeMap<String, IProtoLogGroup> logGroups,
             Runnable cacheUpdater
     ) {
         this(viewerConfigInputStreamProvider,
-                new ProtoLogViewerConfigReader(viewerConfigInputStreamProvider), logGroups,
+                new ProtoLogViewerConfigReader(viewerConfigInputStreamProvider),
                 cacheUpdater);
     }
 
@@ -127,7 +138,6 @@
     public PerfettoProtoLogImpl(
             ViewerConfigInputStreamProvider viewerConfigInputStreamProvider,
             ProtoLogViewerConfigReader viewerConfigReader,
-            TreeMap<String, IProtoLogGroup> logGroups,
             Runnable cacheUpdater
     ) {
         Producer.init(InitArguments.DEFAULTS);
@@ -140,7 +150,6 @@
         mDataSource.register(params);
         this.mViewerConfigInputStreamProvider = viewerConfigInputStreamProvider;
         this.mViewerConfigReader = viewerConfigReader;
-        this.mLogGroups = logGroups;
         this.mCacheUpdater = cacheUpdater;
     }
 
@@ -149,23 +158,70 @@
      */
     @VisibleForTesting
     @Override
-    public void log(LogLevel level, IProtoLogGroup group, long messageHash, int paramsMask,
-            @Nullable String messageString, Object[] args) {
-        Trace.traceBegin(Trace.TRACE_TAG_WINDOW_MANAGER, "log");
+    public void log(LogLevel logLevel, IProtoLogGroup group, long messageHash, int paramsMask,
+            @Nullable Object[] args) {
+        log(logLevel, group, new Message(messageHash, paramsMask), args);
+    }
 
-        long tsNanos = SystemClock.elapsedRealtimeNanos();
-        try {
-            mBackgroundLoggingService.submit(() ->
-                    logToProto(level, group.name(), messageHash, paramsMask, args, tsNanos));
-            if (group.isLogToLogcat()) {
-                logToLogcat(group.getTag(), level, messageHash, messageString, args);
+    @Override
+    public void log(LogLevel logLevel, IProtoLogGroup group, String messageString, Object... args) {
+        log(logLevel, group, new Message(messageString), args);
+    }
+
+    private void log(LogLevel logLevel, IProtoLogGroup group, Message message,
+            @Nullable Object[] args) {
+        if (isProtoEnabled()) {
+            long tsNanos = SystemClock.elapsedRealtimeNanos();
+            final String stacktrace;
+            if (mCollectStackTraceGroupCounts.getOrDefault(group, 0) > 0) {
+                stacktrace = collectStackTrace();
+            } else {
+                stacktrace = null;
             }
-        } finally {
-            Trace.traceEnd(Trace.TRACE_TAG_WINDOW_MANAGER);
+            try {
+                mBackgroundServiceLock.lock();
+                mBackgroundLoggingService.execute(() ->
+                        logToProto(logLevel, group, message, args, tsNanos,
+                                stacktrace));
+            } finally {
+                mBackgroundServiceLock.unlock();
+            }
+        }
+        if (group.isLogToLogcat()) {
+            logToLogcat(group.getTag(), logLevel, message, args);
         }
     }
 
+    private void onTracingFlush() {
+        final ExecutorService loggingService;
+        try {
+            mBackgroundServiceLock.lock();
+            loggingService = mBackgroundLoggingService;
+            mBackgroundLoggingService = Executors.newSingleThreadExecutor();
+        } finally {
+            mBackgroundServiceLock.unlock();
+        }
+
+        try {
+            loggingService.shutdown();
+            boolean finished = loggingService.awaitTermination(10, TimeUnit.SECONDS);
+
+            if (!finished) {
+                Log.e(LOG_TAG, "ProtoLog background tracing service didn't finish gracefully.");
+            }
+        } catch (InterruptedException e) {
+            Log.e(LOG_TAG, "Failed to wait for tracing to finish", e);
+        }
+
+        dumpTransitionTraceConfig();
+    }
+
     private void dumpTransitionTraceConfig() {
+        if (mViewerConfigInputStreamProvider == null) {
+            // No viewer config available
+            return;
+        }
+
         ProtoInputStream pis = mViewerConfigInputStreamProvider.getInputStream();
 
         if (pis == null) {
@@ -256,39 +312,44 @@
         os.end(outMessagesToken);
     }
 
-    private void logToLogcat(String tag, LogLevel level, long messageHash,
-            @Nullable String messageString, Object[] args) {
-        Trace.traceBegin(Trace.TRACE_TAG_WINDOW_MANAGER, "logToLogcat");
-        try {
-            doLogToLogcat(tag, level, messageHash, messageString, args);
-        } finally {
-            Trace.traceEnd(Trace.TRACE_TAG_WINDOW_MANAGER);
+    private void logToLogcat(String tag, LogLevel level, Message message,
+            @Nullable Object[] args) {
+        String messageString = message.getMessage(mViewerConfigReader);
+
+        if (messageString == null) {
+            StringBuilder builder = new StringBuilder("UNKNOWN MESSAGE");
+            if (args != null) {
+                builder.append(" args = (");
+                builder.append(String.join(", ", Arrays.stream(args)
+                        .map(it -> {
+                            if (it == null) {
+                                return "null";
+                            } else {
+                                return it.toString();
+                            }
+                        }).toList()));
+                builder.append(")");
+            }
+            messageString = builder.toString();
+            args = new Object[0];
         }
+
+        logToLogcat(tag, level, messageString, args);
     }
 
-    private void doLogToLogcat(String tag, LogLevel level, long messageHash,
-            @androidx.annotation.Nullable String messageString, Object[] args) {
-        String message = null;
-        if (messageString == null) {
-            messageString = mViewerConfigReader.getViewerString(messageHash);
-        }
-        if (messageString != null) {
-            if (args != null) {
-                try {
-                    message = TextUtils.formatSimple(messageString, args);
-                } catch (Exception ex) {
-                    Slog.w(LOG_TAG, "Invalid ProtoLog format string.", ex);
-                }
-            } else {
-                message = messageString;
+    private void logToLogcat(String tag, LogLevel level, String messageString,
+            @Nullable Object[] args) {
+        String message;
+        if (args != null) {
+            try {
+                message = TextUtils.formatSimple(messageString, args);
+            } catch (IllegalArgumentException e) {
+                message = "FORMAT_ERROR \"" + messageString + "\", args=("
+                        + String.join(
+                                ", ", Arrays.stream(args).map(Object::toString).toList()) + ")";
             }
-        }
-        if (message == null) {
-            StringBuilder builder = new StringBuilder("UNKNOWN MESSAGE (" + messageHash + ")");
-            for (Object o : args) {
-                builder.append(" ").append(o);
-            }
-            message = builder.toString();
+        } else {
+            message = messageString;
         }
         passToLogcat(tag, level, message);
     }
@@ -320,25 +381,11 @@
         }
     }
 
-    private void logToProto(LogLevel level, String groupName, long messageHash, int paramsMask,
-            Object[] args, long tsNanos) {
-        if (!isProtoEnabled()) {
-            return;
-        }
-
-        Trace.traceBegin(Trace.TRACE_TAG_WINDOW_MANAGER, "logToProto");
-        try {
-            doLogToProto(level, groupName, messageHash, paramsMask, args, tsNanos);
-        } finally {
-            Trace.traceEnd(Trace.TRACE_TAG_WINDOW_MANAGER);
-        }
-    }
-
-    private void doLogToProto(LogLevel level, String groupName, long messageHash, int paramsMask,
-            Object[] args, long tsNanos) {
+    private void logToProto(LogLevel level, IProtoLogGroup logGroup, Message message, Object[] args,
+            long tsNanos, @Nullable String stacktrace) {
         mDataSource.trace(ctx -> {
             final ProtoLogDataSource.TlsState tlsState = ctx.getCustomTlsState();
-            final LogLevel logFrom = tlsState.getLogFromLevel(groupName);
+            final LogLevel logFrom = tlsState.getLogFromLevel(logGroup.name());
 
             if (level.ordinal() < logFrom.ordinal()) {
                 return;
@@ -350,29 +397,43 @@
                 // trace processing easier.
                 int argIndex = 0;
                 for (Object o : args) {
-                    int type = LogDataType.bitmaskToLogDataType(paramsMask, argIndex);
+                    int type = LogDataType.bitmaskToLogDataType(message.getMessageMask(), argIndex);
                     if (type == LogDataType.STRING) {
-                        internStringArg(ctx, o.toString());
+                        if (o == null) {
+                            internStringArg(ctx, NULL_STRING);
+                        } else {
+                            internStringArg(ctx, o.toString());
+                        }
                     }
                     argIndex++;
                 }
             }
 
             int internedStacktrace = 0;
-            if (tlsState.getShouldCollectStacktrace(groupName)) {
+            if (tlsState.getShouldCollectStacktrace(logGroup.name())) {
                 // Intern stackstraces before creating the trace packet for the proto message so
                 // that the interned stacktrace strings appear before in the trace to make the
                 // trace processing easier.
-                String stacktrace = collectStackTrace();
                 internedStacktrace = internStacktraceString(ctx, stacktrace);
             }
 
+            boolean needsIncrementalState = false;
+
+            long messageHash = 0;
+            if (message.mMessageHash != null) {
+                messageHash = message.mMessageHash;
+            }
+            if (message.mMessageString != null) {
+                needsIncrementalState = true;
+                messageHash =
+                        internProtoMessage(ctx, level, logGroup, message.mMessageString);
+            }
+
             final ProtoOutputStream os = ctx.newTracePacket();
             os.write(TIMESTAMP, tsNanos);
             long token = os.start(PROTOLOG_MESSAGE);
-            os.write(MESSAGE_ID, messageHash);
 
-            boolean needsIncrementalState = false;
+            os.write(MESSAGE_ID, messageHash);
 
             if (args != null) {
 
@@ -381,22 +442,39 @@
                 ArrayList<Double> doubleParams = new ArrayList<>();
                 ArrayList<Boolean> booleanParams = new ArrayList<>();
                 for (Object o : args) {
-                    int type = LogDataType.bitmaskToLogDataType(paramsMask, argIndex);
+                    int type = LogDataType.bitmaskToLogDataType(message.getMessageMask(), argIndex);
                     try {
                         switch (type) {
                             case LogDataType.STRING:
-                                final int internedStringId = internStringArg(ctx, o.toString());
+                                final int internedStringId;
+                                if (o == null) {
+                                    internedStringId = internStringArg(ctx, NULL_STRING);
+                                } else {
+                                    internedStringId = internStringArg(ctx, o.toString());
+                                }
                                 os.write(STR_PARAM_IIDS, internedStringId);
                                 needsIncrementalState = true;
                                 break;
                             case LogDataType.LONG:
-                                longParams.add(((Number) o).longValue());
+                                if (o == null) {
+                                    longParams.add(0);
+                                } else {
+                                    longParams.add(((Number) o).longValue());
+                                }
                                 break;
                             case LogDataType.DOUBLE:
-                                doubleParams.add(((Number) o).doubleValue());
+                                if (o == null) {
+                                    doubleParams.add(0d);
+                                } else {
+                                    doubleParams.add(((Number) o).doubleValue());
+                                }
                                 break;
                             case LogDataType.BOOLEAN:
-                                booleanParams.add((boolean) o);
+                                if (o == null) {
+                                    booleanParams.add(false);
+                                } else {
+                                    booleanParams.add((boolean) o);
+                                }
                                 break;
                         }
                     } catch (ClassCastException ex) {
@@ -414,7 +492,7 @@
                 booleanParams.forEach(it -> os.write(BOOLEAN_PARAMS, it ? 1 : 0));
             }
 
-            if (tlsState.getShouldCollectStacktrace(groupName)) {
+            if (tlsState.getShouldCollectStacktrace(logGroup.name())) {
                 os.write(STACKTRACE_IID, internedStacktrace);
             }
 
@@ -427,6 +505,63 @@
         });
     }
 
+    private long internProtoMessage(
+            TracingContext<ProtoLogDataSource.Instance, ProtoLogDataSource.TlsState,
+                    ProtoLogDataSource.IncrementalState> ctx, LogLevel level,
+            IProtoLogGroup logGroup, String message) {
+        final ProtoLogDataSource.IncrementalState incrementalState = ctx.getIncrementalState();
+
+        if (!incrementalState.clearReported) {
+            final ProtoOutputStream os = ctx.newTracePacket();
+            os.write(SEQUENCE_FLAGS, SEQ_INCREMENTAL_STATE_CLEARED);
+            incrementalState.clearReported = true;
+        }
+
+
+        if (!incrementalState.protologGroupInterningSet.contains(logGroup.getId())) {
+            incrementalState.protologGroupInterningSet.add(logGroup.getId());
+
+            final ProtoOutputStream os = ctx.newTracePacket();
+            final long protologViewerConfigToken = os.start(PROTOLOG_VIEWER_CONFIG);
+            final long groupConfigToken = os.start(GROUPS);
+
+            os.write(ID, logGroup.getId());
+            os.write(NAME, logGroup.name());
+            os.write(TAG, logGroup.getTag());
+
+            os.end(groupConfigToken);
+            os.end(protologViewerConfigToken);
+        }
+
+        final Long messageHash = hash(level, logGroup.name(), message);
+        if (!incrementalState.protologMessageInterningSet.contains(messageHash)) {
+            incrementalState.protologMessageInterningSet.add(messageHash);
+
+            final ProtoOutputStream os = ctx.newTracePacket();
+            final long protologViewerConfigToken = os.start(PROTOLOG_VIEWER_CONFIG);
+            final long messageConfigToken = os.start(MESSAGES);
+
+            os.write(MessageData.MESSAGE_ID, messageHash);
+            os.write(MESSAGE, message);
+            os.write(LEVEL, level.ordinal());
+            os.write(GROUP_ID, logGroup.getId());
+
+            os.end(messageConfigToken);
+            os.end(protologViewerConfigToken);
+        }
+
+        return messageHash;
+    }
+
+    private Long hash(
+            LogLevel logLevel,
+            String logGroup,
+            String messageString
+    ) {
+        final String fullStringIdentifier =  messageString + logLevel + logGroup;
+        return UUID.nameUUIDFromBytes(fullStringIdentifier.getBytes()).getMostSignificantBits();
+    }
+
     private static final int STACK_SIZE_TO_PROTO_LOG_ENTRY_CALL = 12;
 
     private String collectStackTrace() {
@@ -466,7 +601,7 @@
                     ProtoLogDataSource.IncrementalState> ctx,
             Map<String, Integer> internMap,
             long fieldId,
-            String string
+            @NonNull String string
     ) {
         final ProtoLogDataSource.IncrementalState incrementalState = ctx.getIncrementalState();
 
@@ -523,25 +658,17 @@
 
     @Override
     public boolean isEnabled(IProtoLogGroup group, LogLevel level) {
-        return group.isLogToLogcat() || getLogFromLevel(group).ordinal() <= level.ordinal();
+        final int[] groupLevelCount = mLogLevelCounts.get(group);
+        return (groupLevelCount == null && mDefaultLogLevelCounts[level.ordinal()] > 0)
+                || (groupLevelCount != null && groupLevelCount[level.ordinal()] > 0)
+                || group.isLogToLogcat();
     }
 
-    private LogLevel getLogFromLevel(IProtoLogGroup group) {
-        if (mLogLevelCounts.containsKey(group)) {
-            for (LogLevel logLevel : LogLevel.values()) {
-                if (mLogLevelCounts.get(group).getOrDefault(logLevel, 0) > 0) {
-                    return logLevel;
-                }
-            }
-        } else {
-            for (LogLevel logLevel : LogLevel.values()) {
-                if (mDefaultLogLevelCounts.getOrDefault(logLevel, 0) > 0) {
-                    return logLevel;
-                }
-            }
+    @Override
+    public void registerGroups(IProtoLogGroup... protoLogGroups) {
+        for (IProtoLogGroup protoLogGroup : protoLogGroups) {
+            mLogGroups.put(protoLogGroup.name(), protoLogGroup);
         }
-
-        return LogLevel.WTF;
     }
 
     /**
@@ -620,36 +747,51 @@
     }
 
     private synchronized void onTracingInstanceStart(ProtoLogDataSource.ProtoLogConfig config) {
-        this.mTracingInstances.incrementAndGet();
-
         final LogLevel defaultLogFrom = config.getDefaultGroupConfig().logFrom;
-        mDefaultLogLevelCounts.put(defaultLogFrom,
-                mDefaultLogLevelCounts.getOrDefault(defaultLogFrom, 0) + 1);
+        for (int i = defaultLogFrom.ordinal(); i < LogLevel.values().length; i++) {
+            mDefaultLogLevelCounts[i]++;
+        }
 
         final Set<String> overriddenGroupTags = config.getGroupTagsWithOverriddenConfigs();
 
         for (String overriddenGroupTag : overriddenGroupTags) {
             IProtoLogGroup group = mLogGroups.get(overriddenGroupTag);
 
-            mLogLevelCounts.putIfAbsent(group, new ArrayMap<>());
-            final Map<LogLevel, Integer> logLevelsCountsForGroup = mLogLevelCounts.get(group);
+            if (group == null) {
+                throw new IllegalArgumentException("Trying to set config for \""
+                        + overriddenGroupTag + "\" that isn't registered");
+            }
+
+            mLogLevelCounts.putIfAbsent(group, new int[LogLevel.values().length]);
+            final int[] logLevelsCountsForGroup = mLogLevelCounts.get(group);
 
             final LogLevel logFromLevel = config.getConfigFor(overriddenGroupTag).logFrom;
-            logLevelsCountsForGroup.put(logFromLevel,
-                    logLevelsCountsForGroup.getOrDefault(logFromLevel, 0) + 1);
+            for (int i = logFromLevel.ordinal(); i < LogLevel.values().length; i++) {
+                logLevelsCountsForGroup[i]++;
+            }
+
+            if (config.getConfigFor(overriddenGroupTag).collectStackTrace) {
+                mCollectStackTraceGroupCounts.put(group,
+                        mCollectStackTraceGroupCounts.getOrDefault(group, 0) + 1);
+            }
+
+            if (config.getConfigFor(overriddenGroupTag).collectStackTrace) {
+                mCollectStackTraceGroupCounts.put(group,
+                        mCollectStackTraceGroupCounts.getOrDefault(group, 0) + 1);
+            }
         }
 
         mCacheUpdater.run();
+
+        this.mTracingInstances.incrementAndGet();
     }
 
     private synchronized void onTracingInstanceStop(ProtoLogDataSource.ProtoLogConfig config) {
         this.mTracingInstances.decrementAndGet();
 
         final LogLevel defaultLogFrom = config.getDefaultGroupConfig().logFrom;
-        mDefaultLogLevelCounts.put(defaultLogFrom,
-                mDefaultLogLevelCounts.get(defaultLogFrom) - 1);
-        if (mDefaultLogLevelCounts.get(defaultLogFrom) <= 0) {
-            mDefaultLogLevelCounts.remove(defaultLogFrom);
+        for (int i = defaultLogFrom.ordinal(); i < LogLevel.values().length; i++) {
+            mDefaultLogLevelCounts[i]--;
         }
 
         final Set<String> overriddenGroupTags = config.getGroupTagsWithOverriddenConfigs();
@@ -657,18 +799,24 @@
         for (String overriddenGroupTag : overriddenGroupTags) {
             IProtoLogGroup group = mLogGroups.get(overriddenGroupTag);
 
-            mLogLevelCounts.putIfAbsent(group, new ArrayMap<>());
-            final Map<LogLevel, Integer> logLevelsCountsForGroup = mLogLevelCounts.get(group);
+            final int[] logLevelsCountsForGroup = mLogLevelCounts.get(group);
 
             final LogLevel logFromLevel = config.getConfigFor(overriddenGroupTag).logFrom;
-            logLevelsCountsForGroup.put(logFromLevel,
-                    logLevelsCountsForGroup.get(logFromLevel) - 1);
-            if (logLevelsCountsForGroup.get(logFromLevel) <= 0) {
-                logLevelsCountsForGroup.remove(logFromLevel);
+            for (int i = defaultLogFrom.ordinal(); i < LogLevel.values().length; i++) {
+                logLevelsCountsForGroup[i]--;
             }
-            if (logLevelsCountsForGroup.isEmpty()) {
+            if (Arrays.stream(logLevelsCountsForGroup).allMatch(it -> it == 0)) {
                 mLogLevelCounts.remove(group);
             }
+
+            if (config.getConfigFor(overriddenGroupTag).collectStackTrace) {
+                mCollectStackTraceGroupCounts.put(group,
+                        mCollectStackTraceGroupCounts.get(group) - 1);
+
+                if (mCollectStackTraceGroupCounts.get(group) == 0) {
+                    mCollectStackTraceGroupCounts.remove(group);
+                }
+            }
         }
 
         mCacheUpdater.run();
@@ -681,5 +829,36 @@
             pw.flush();
         }
     }
+
+    private static class Message {
+        private final Long mMessageHash;
+        private final Integer mMessageMask;
+        private final String mMessageString;
+
+        private Message(Long messageHash, int messageMask) {
+            this.mMessageHash = messageHash;
+            this.mMessageMask = messageMask;
+            this.mMessageString = null;
+        }
+
+        private Message(String messageString) {
+            this.mMessageHash = null;
+            final List<Integer> argTypes = LogDataType.parseFormatString(messageString);
+            this.mMessageMask = LogDataType.logDataTypesToBitMask(argTypes);
+            this.mMessageString = messageString;
+        }
+
+        private int getMessageMask() {
+            return mMessageMask;
+        }
+
+        private String getMessage(ProtoLogViewerConfigReader viewerConfigReader) {
+            if (mMessageString != null) {
+                return mMessageString;
+            }
+
+            return viewerConfigReader.getViewerString(mMessageHash);
+        }
+    }
 }
 
diff --git a/core/java/com/android/internal/protolog/ProtoLog.java b/core/java/com/android/internal/protolog/ProtoLog.java
index 0118c05..87678e5 100644
--- a/core/java/com/android/internal/protolog/ProtoLog.java
+++ b/core/java/com/android/internal/protolog/ProtoLog.java
@@ -44,21 +44,23 @@
 // LINT.ThenChange(frameworks/base/tools/protologtool/src/com/android/protolog/tool/ProtoLogTool.kt)
 
     // Needs to be set directly otherwise the protologtool tries to transform the method call
+    @Deprecated
     public static boolean REQUIRE_PROTOLOGTOOL = true;
 
+    private static IProtoLog sProtoLogInstance;
+
     /**
      * DEBUG level log.
      *
      * @param group         {@code IProtoLogGroup} controlling this log call.
      * @param messageString constant format string for the logged message.
      * @param args          parameters to be used with the format string.
+     *
+     * NOTE: If source code is pre-processed by ProtoLogTool this is not the function call that is
+     *       executed. Check generated code for actual call.
      */
     public static void d(IProtoLogGroup group, String messageString, Object... args) {
-        // Stub, replaced by the ProtoLogTool.
-        if (REQUIRE_PROTOLOGTOOL) {
-            throw new UnsupportedOperationException(
-                    "ProtoLog calls MUST be processed with ProtoLogTool");
-        }
+        logStringMessage(LogLevel.DEBUG, group, messageString, args);
     }
 
     /**
@@ -67,13 +69,12 @@
      * @param group         {@code IProtoLogGroup} controlling this log call.
      * @param messageString constant format string for the logged message.
      * @param args          parameters to be used with the format string.
+     *
+     * NOTE: If source code is pre-processed by ProtoLogTool this is not the function call that is
+     *       executed. Check generated code for actual call.
      */
     public static void v(IProtoLogGroup group, String messageString, Object... args) {
-        // Stub, replaced by the ProtoLogTool.
-        if (REQUIRE_PROTOLOGTOOL) {
-            throw new UnsupportedOperationException(
-                    "ProtoLog calls MUST be processed with ProtoLogTool");
-        }
+        logStringMessage(LogLevel.VERBOSE, group, messageString, args);
     }
 
     /**
@@ -82,13 +83,12 @@
      * @param group         {@code IProtoLogGroup} controlling this log call.
      * @param messageString constant format string for the logged message.
      * @param args          parameters to be used with the format string.
+     *
+     * NOTE: If source code is pre-processed by ProtoLogTool this is not the function call that is
+     *       executed. Check generated code for actual call.
      */
     public static void i(IProtoLogGroup group, String messageString, Object... args) {
-        // Stub, replaced by the ProtoLogTool.
-        if (REQUIRE_PROTOLOGTOOL) {
-            throw new UnsupportedOperationException(
-                    "ProtoLog calls MUST be processed with ProtoLogTool");
-        }
+        logStringMessage(LogLevel.INFO, group, messageString, args);
     }
 
     /**
@@ -97,13 +97,12 @@
      * @param group         {@code IProtoLogGroup} controlling this log call.
      * @param messageString constant format string for the logged message.
      * @param args          parameters to be used with the format string.
+     *
+     * NOTE: If source code is pre-processed by ProtoLogTool this is not the function call that is
+     *       executed. Check generated code for actual call.
      */
     public static void w(IProtoLogGroup group, String messageString, Object... args) {
-        // Stub, replaced by the ProtoLogTool.
-        if (REQUIRE_PROTOLOGTOOL) {
-            throw new UnsupportedOperationException(
-                    "ProtoLog calls MUST be processed with ProtoLogTool");
-        }
+        logStringMessage(LogLevel.WARN, group, messageString, args);
     }
 
     /**
@@ -112,13 +111,12 @@
      * @param group         {@code IProtoLogGroup} controlling this log call.
      * @param messageString constant format string for the logged message.
      * @param args          parameters to be used with the format string.
+     *
+     * NOTE: If source code is pre-processed by ProtoLogTool this is not the function call that is
+     *       executed. Check generated code for actual call.
      */
     public static void e(IProtoLogGroup group, String messageString, Object... args) {
-        // Stub, replaced by the ProtoLogTool.
-        if (REQUIRE_PROTOLOGTOOL) {
-            throw new UnsupportedOperationException(
-                    "ProtoLog calls MUST be processed with ProtoLogTool");
-        }
+        logStringMessage(LogLevel.ERROR, group, messageString, args);
     }
 
     /**
@@ -127,13 +125,12 @@
      * @param group         {@code IProtoLogGroup} controlling this log call.
      * @param messageString constant format string for the logged message.
      * @param args          parameters to be used with the format string.
+     *
+     * NOTE: If source code is pre-processed by ProtoLogTool this is not the function call that is
+     *       executed. Check generated code for actual call.
      */
     public static void wtf(IProtoLogGroup group, String messageString, Object... args) {
-        // Stub, replaced by the ProtoLogTool.
-        if (REQUIRE_PROTOLOGTOOL) {
-            throw new UnsupportedOperationException(
-                    "ProtoLog calls MUST be processed with ProtoLogTool");
-        }
+        logStringMessage(LogLevel.WTF, group, messageString, args);
     }
 
     /**
@@ -142,11 +139,7 @@
      * @return true iff this is being logged.
      */
     public static boolean isEnabled(IProtoLogGroup group, LogLevel level) {
-        if (REQUIRE_PROTOLOGTOOL) {
-            throw new UnsupportedOperationException(
-                    "ProtoLog calls MUST be processed with ProtoLogTool");
-        }
-        return false;
+        return sProtoLogInstance.isEnabled(group, level);
     }
 
     /**
@@ -154,10 +147,36 @@
      * @return A singleton instance of ProtoLog.
      */
     public static IProtoLog getSingleInstance() {
-        if (REQUIRE_PROTOLOGTOOL) {
-            throw new UnsupportedOperationException(
-                    "ProtoLog calls MUST be processed with ProtoLogTool");
+        return sProtoLogInstance;
+    }
+
+    /**
+     * Registers available protolog groups. A group must be registered before it can be used.
+     * @param protoLogGroups The groups to register for use in protolog.
+     */
+    public static void registerGroups(IProtoLogGroup... protoLogGroups) {
+        sProtoLogInstance.registerGroups(protoLogGroups);
+    }
+
+    private static void logStringMessage(LogLevel logLevel, IProtoLogGroup group,
+            String stringMessage, Object... args) {
+        if (sProtoLogInstance == null) {
+            throw new IllegalStateException(
+                    "Trying to use ProtoLog before it is initialized in this process.");
         }
-        return null;
+
+        if (sProtoLogInstance.isEnabled(group, logLevel)) {
+            sProtoLogInstance.log(logLevel, group, stringMessage, args);
+        }
+    }
+
+    static {
+        if (android.tracing.Flags.perfettoProtologTracing()) {
+            sProtoLogInstance = new PerfettoProtoLogImpl();
+        } else {
+            // The first call to ProtoLog is likely to flip REQUIRE_PROTOLOGTOOL, which is when this
+            // static block will be executed before REQUIRE_PROTOLOGTOOL is actually set.
+            sProtoLogInstance = new LogcatOnlyProtoLogImpl();
+        }
     }
 }
diff --git a/core/java/com/android/internal/protolog/ProtoLogDataSource.java b/core/java/com/android/internal/protolog/ProtoLogDataSource.java
index 6ab79b9..84f3237 100644
--- a/core/java/com/android/internal/protolog/ProtoLogDataSource.java
+++ b/core/java/com/android/internal/protolog/ProtoLogDataSource.java
@@ -40,6 +40,7 @@
 
 import java.io.IOException;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.function.Consumer;
@@ -138,6 +139,8 @@
     }
 
     public static class IncrementalState {
+        public final Set<Integer> protologGroupInterningSet = new HashSet<>();
+        public final Set<Long> protologMessageInterningSet = new HashSet<>();
         public final Map<String, Integer> argumentInterningMap = new HashMap<>();
         public final Map<String, Integer> stacktraceInterningMap = new HashMap<>();
         public boolean clearReported = false;
diff --git a/core/java/com/android/internal/protolog/ProtoLogImpl.java b/core/java/com/android/internal/protolog/ProtoLogImpl.java
index 6d142af..3082295 100644
--- a/core/java/com/android/internal/protolog/ProtoLogImpl.java
+++ b/core/java/com/android/internal/protolog/ProtoLogImpl.java
@@ -54,48 +54,33 @@
     private static Runnable sCacheUpdater;
 
     /** Used by the ProtoLogTool, do not call directly - use {@code ProtoLog} class instead. */
-    public static void d(IProtoLogGroup group, long messageHash, int paramsMask,
-            @Nullable String messageString,
-            Object... args) {
-        getSingleInstance()
-                .log(LogLevel.DEBUG, group, messageHash, paramsMask, messageString, args);
+    public static void d(IProtoLogGroup group, long messageHash, int paramsMask, Object... args) {
+        getSingleInstance().log(LogLevel.DEBUG, group, messageHash, paramsMask, args);
     }
 
     /** Used by the ProtoLogTool, do not call directly - use {@code ProtoLog} class instead. */
-    public static void v(IProtoLogGroup group, long messageHash, int paramsMask,
-            @Nullable String messageString,
-            Object... args) {
-        getSingleInstance().log(LogLevel.VERBOSE, group, messageHash, paramsMask, messageString,
-                args);
+    public static void v(IProtoLogGroup group, long messageHash, int paramsMask, Object... args) {
+        getSingleInstance().log(LogLevel.VERBOSE, group, messageHash, paramsMask, args);
     }
 
     /** Used by the ProtoLogTool, do not call directly - use {@code ProtoLog} class instead. */
-    public static void i(IProtoLogGroup group, long messageHash, int paramsMask,
-            @Nullable String messageString,
-            Object... args) {
-        getSingleInstance().log(LogLevel.INFO, group, messageHash, paramsMask, messageString, args);
+    public static void i(IProtoLogGroup group, long messageHash, int paramsMask, Object... args) {
+        getSingleInstance().log(LogLevel.INFO, group, messageHash, paramsMask, args);
     }
 
     /** Used by the ProtoLogTool, do not call directly - use {@code ProtoLog} class instead. */
-    public static void w(IProtoLogGroup group, long messageHash, int paramsMask,
-            @Nullable String messageString,
-            Object... args) {
-        getSingleInstance().log(LogLevel.WARN, group, messageHash, paramsMask, messageString, args);
+    public static void w(IProtoLogGroup group, long messageHash, int paramsMask, Object... args) {
+        getSingleInstance().log(LogLevel.WARN, group, messageHash, paramsMask, args);
     }
 
     /** Used by the ProtoLogTool, do not call directly - use {@code ProtoLog} class instead. */
-    public static void e(IProtoLogGroup group, long messageHash, int paramsMask,
-            @Nullable String messageString,
-            Object... args) {
-        getSingleInstance()
-                .log(LogLevel.ERROR, group, messageHash, paramsMask, messageString, args);
+    public static void e(IProtoLogGroup group, long messageHash, int paramsMask, Object... args) {
+        getSingleInstance().log(LogLevel.ERROR, group, messageHash, paramsMask, args);
     }
 
     /** Used by the ProtoLogTool, do not call directly - use {@code ProtoLog} class instead. */
-    public static void wtf(IProtoLogGroup group, long messageHash, int paramsMask,
-            @Nullable String messageString,
-            Object... args) {
-        getSingleInstance().log(LogLevel.WTF, group, messageHash, paramsMask, messageString, args);
+    public static void wtf(IProtoLogGroup group, long messageHash, int paramsMask, Object... args) {
+        getSingleInstance().log(LogLevel.WTF, group, messageHash, paramsMask, args);
     }
 
     /**
@@ -107,18 +92,27 @@
     }
 
     /**
+     * Registers available protolog groups. A group must be registered before it can be used.
+     * @param protoLogGroups The groups to register for use in protolog.
+     */
+    public static void registerGroups(IProtoLogGroup... protoLogGroups) {
+        getSingleInstance().registerGroups(protoLogGroups);
+    }
+
+    /**
      * Returns the single instance of the ProtoLogImpl singleton class.
      */
     public static synchronized IProtoLog getSingleInstance() {
         if (sServiceInstance == null) {
             if (android.tracing.Flags.perfettoProtologTracing()) {
-                sServiceInstance = new PerfettoProtoLogImpl(
-                        sViewerConfigPath, sLogGroups, sCacheUpdater);
+                sServiceInstance = new PerfettoProtoLogImpl(sViewerConfigPath, sCacheUpdater);
             } else {
                 sServiceInstance = new LegacyProtoLogImpl(
-                        sLegacyOutputFilePath, sLegacyViewerConfigPath, sLogGroups, sCacheUpdater);
+                        sLegacyOutputFilePath, sLegacyViewerConfigPath, sCacheUpdater);
             }
 
+            IProtoLogGroup[] groups = sLogGroups.values().toArray(new IProtoLogGroup[0]);
+            sServiceInstance.registerGroups(groups);
             sCacheUpdater.run();
         }
         return sServiceInstance;
diff --git a/core/java/com/android/internal/protolog/common/IProtoLog.java b/core/java/com/android/internal/protolog/common/IProtoLog.java
index f72d9f7..f5695ac 100644
--- a/core/java/com/android/internal/protolog/common/IProtoLog.java
+++ b/core/java/com/android/internal/protolog/common/IProtoLog.java
@@ -27,11 +27,19 @@
      * @param group The group this message belongs to.
      * @param messageHash The hash of the message.
      * @param paramsMask The parameters mask of the message.
-     * @param messageString The message string.
      * @param args The arguments of the message.
      */
     void log(LogLevel logLevel, IProtoLogGroup group, long messageHash, int paramsMask,
-             String messageString, Object[] args);
+            Object[] args);
+
+    /**
+     * Log a ProtoLog message
+     * @param logLevel Log level of the proto message.
+     * @param group The group this message belongs to.
+     * @param messageString The message string.
+     * @param args The arguments of the message.
+     */
+    void log(LogLevel logLevel, IProtoLogGroup group, String messageString, Object... args);
 
     /**
      * Check if ProtoLog is tracing.
@@ -60,4 +68,10 @@
      * @return If we need to log this group and level to either ProtoLog or Logcat.
      */
     boolean isEnabled(IProtoLogGroup group, LogLevel level);
+
+    /**
+     * Registers available protolog groups. A group must be registered before it can be used.
+     * @param protoLogGroups The groups to register for use in protolog.
+     */
+    void registerGroups(IProtoLogGroup... protoLogGroups);
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/shortcut/CreateBubbleShortcutActivity.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/shortcut/CreateBubbleShortcutActivity.kt
index a124f95..c93c11e 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/shortcut/CreateBubbleShortcutActivity.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/shortcut/CreateBubbleShortcutActivity.kt
@@ -20,10 +20,10 @@
 import android.content.pm.ShortcutManager
 import android.graphics.drawable.Icon
 import android.os.Bundle
+import com.android.internal.protolog.ProtoLog
 import com.android.wm.shell.Flags
 import com.android.wm.shell.R
 import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_BUBBLES
-import com.android.wm.shell.util.KtProtoLog
 
 /** Activity to create a shortcut to open bubbles */
 class CreateBubbleShortcutActivity : Activity() {
@@ -31,7 +31,7 @@
     override fun onCreate(savedInstanceState: Bundle?) {
         super.onCreate(savedInstanceState)
         if (Flags.enableRetrievableBubbles()) {
-            KtProtoLog.d(WM_SHELL_BUBBLES, "Creating a shortcut for bubbles")
+            ProtoLog.d(WM_SHELL_BUBBLES, "Creating a shortcut for bubbles")
             createShortcut()
         }
         finish()
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/shortcut/ShowBubblesActivity.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/shortcut/ShowBubblesActivity.kt
index ae7940c..e578e9e 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/shortcut/ShowBubblesActivity.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/shortcut/ShowBubblesActivity.kt
@@ -21,9 +21,9 @@
 import android.content.Context
 import android.content.Intent
 import android.os.Bundle
+import com.android.internal.protolog.ProtoLog
 import com.android.wm.shell.Flags
 import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_BUBBLES
-import com.android.wm.shell.util.KtProtoLog
 
 /** Activity that sends a broadcast to open bubbles */
 class ShowBubblesActivity : Activity() {
@@ -37,7 +37,7 @@
                     // Set the package as the receiver is not exported
                     `package` = packageName
                 }
-            KtProtoLog.v(WM_SHELL_BUBBLES, "Sending broadcast to show bubbles")
+            ProtoLog.v(WM_SHELL_BUBBLES, "Sending broadcast to show bubbles")
             sendBroadcast(intent)
         }
         finish()
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/LaunchAdjacentController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/common/LaunchAdjacentController.kt
index 81592c3..e92b0b5 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/LaunchAdjacentController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/LaunchAdjacentController.kt
@@ -17,8 +17,8 @@
 
 import android.window.WindowContainerToken
 import android.window.WindowContainerTransaction
+import com.android.internal.protolog.ProtoLog
 import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_TASK_ORG
-import com.android.wm.shell.util.KtProtoLog
 
 /**
  * Controller to manage behavior of activities launched with
@@ -30,7 +30,7 @@
     var launchAdjacentEnabled: Boolean = true
         set(value) {
             if (field != value) {
-                KtProtoLog.d(WM_SHELL_TASK_ORG, "set launch adjacent flag root enabled=%b", value)
+                ProtoLog.d(WM_SHELL_TASK_ORG, "set launch adjacent flag root enabled=%b", value)
                 field = value
                 container?.let { c ->
                     if (value) {
@@ -52,7 +52,7 @@
      * @see WindowContainerTransaction.setLaunchAdjacentFlagRoot
      */
     fun setLaunchAdjacentRoot(container: WindowContainerToken) {
-        KtProtoLog.d(WM_SHELL_TASK_ORG, "set new launch adjacent flag root container")
+        ProtoLog.d(WM_SHELL_TASK_ORG, "set new launch adjacent flag root container")
         this.container = container
         if (launchAdjacentEnabled) {
             enableContainer(container)
@@ -67,7 +67,7 @@
      * @see WindowContainerTransaction.clearLaunchAdjacentFlagRoot
      */
     fun clearLaunchAdjacentRoot() {
-        KtProtoLog.d(WM_SHELL_TASK_ORG, "clear launch adjacent flag root container")
+        ProtoLog.d(WM_SHELL_TASK_ORG, "clear launch adjacent flag root container")
         container?.let {
             disableContainer(it)
             container = null
@@ -75,14 +75,14 @@
     }
 
     private fun enableContainer(container: WindowContainerToken) {
-        KtProtoLog.v(WM_SHELL_TASK_ORG, "enable launch adjacent flag root container")
+        ProtoLog.v(WM_SHELL_TASK_ORG, "enable launch adjacent flag root container")
         val wct = WindowContainerTransaction()
         wct.setLaunchAdjacentFlagRoot(container)
         syncQueue.queue(wct)
     }
 
     private fun disableContainer(container: WindowContainerToken) {
-        KtProtoLog.v(WM_SHELL_TASK_ORG, "disable launch adjacent flag root container")
+        ProtoLog.v(WM_SHELL_TASK_ORG, "disable launch adjacent flag root container")
         val wct = WindowContainerTransaction()
         wct.clearLaunchAdjacentFlagRoot(container)
         syncQueue.queue(wct)
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/MultiInstanceHelper.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/common/MultiInstanceHelper.kt
index 9e8dfb5..a6be640 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/MultiInstanceHelper.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/MultiInstanceHelper.kt
@@ -23,9 +23,9 @@
 import android.os.UserHandle
 import android.view.WindowManager.PROPERTY_SUPPORTS_MULTI_INSTANCE_SYSTEM_UI
 import com.android.internal.annotations.VisibleForTesting
+import com.android.internal.protolog.ProtoLog
 import com.android.wm.shell.R
 import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL
-import com.android.wm.shell.util.KtProtoLog
 import java.util.Arrays
 
 /**
@@ -52,7 +52,7 @@
         val packageName = componentName.packageName
         for (pkg in staticAppsSupportingMultiInstance) {
             if (pkg == packageName) {
-                KtProtoLog.v(WM_SHELL, "application=%s in allowlist supports multi-instance",
+                ProtoLog.v(WM_SHELL, "application=%s in allowlist supports multi-instance",
                     packageName)
                 return true
             }
@@ -70,10 +70,10 @@
             // If the above call doesn't throw a NameNotFoundException, then the activity property
             // should override the application property value
             if (activityProp.isBoolean) {
-                KtProtoLog.v(WM_SHELL, "activity=%s supports multi-instance", componentName)
+                ProtoLog.v(WM_SHELL, "activity=%s supports multi-instance", componentName)
                 return activityProp.boolean
             } else {
-                KtProtoLog.w(WM_SHELL, "Warning: property=%s for activity=%s has non-bool type=%d",
+                ProtoLog.w(WM_SHELL, "Warning: property=%s for activity=%s has non-bool type=%d",
                     PROPERTY_SUPPORTS_MULTI_INSTANCE_SYSTEM_UI, packageName, activityProp.type)
             }
         } catch (nnfe: PackageManager.NameNotFoundException) {
@@ -85,10 +85,10 @@
             val appProp = packageManager.getProperty(
                 PROPERTY_SUPPORTS_MULTI_INSTANCE_SYSTEM_UI, packageName)
             if (appProp.isBoolean) {
-                KtProtoLog.v(WM_SHELL, "application=%s supports multi-instance", packageName)
+                ProtoLog.v(WM_SHELL, "application=%s supports multi-instance", packageName)
                 return appProp.boolean
             } else {
-                KtProtoLog.w(WM_SHELL,
+                ProtoLog.w(WM_SHELL,
                     "Warning: property=%s for application=%s has non-bool type=%d",
                     PROPERTY_SUPPORTS_MULTI_INSTANCE_SYSTEM_UI, packageName, appProp.type)
             }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeEventLogger.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeEventLogger.kt
index fbc11c1..400882a 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeEventLogger.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeEventLogger.kt
@@ -16,9 +16,9 @@
 
 package com.android.wm.shell.desktopmode
 
+import com.android.internal.protolog.ProtoLog
 import com.android.internal.util.FrameworkStatsLog
 import com.android.wm.shell.protolog.ShellProtoLogGroup
-import com.android.wm.shell.util.KtProtoLog
 
 /** Event logger for logging desktop mode session events */
 class DesktopModeEventLogger {
@@ -27,7 +27,7 @@
      * entering desktop mode
      */
     fun logSessionEnter(sessionId: Int, enterReason: EnterReason) {
-        KtProtoLog.v(
+        ProtoLog.v(
             ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
             "DesktopModeLogger: Logging session enter, session: %s reason: %s",
             sessionId,
@@ -47,7 +47,7 @@
      * exiting desktop mode
      */
     fun logSessionExit(sessionId: Int, exitReason: ExitReason) {
-        KtProtoLog.v(
+        ProtoLog.v(
             ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
             "DesktopModeLogger: Logging session exit, session: %s reason: %s",
             sessionId,
@@ -67,7 +67,7 @@
      * session id [sessionId]
      */
     fun logTaskAdded(sessionId: Int, taskUpdate: TaskUpdate) {
-        KtProtoLog.v(
+        ProtoLog.v(
             ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
             "DesktopModeLogger: Logging task added, session: %s taskId: %s",
             sessionId,
@@ -99,7 +99,7 @@
      * session id [sessionId]
      */
     fun logTaskRemoved(sessionId: Int, taskUpdate: TaskUpdate) {
-        KtProtoLog.v(
+        ProtoLog.v(
             ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
             "DesktopModeLogger: Logging task remove, session: %s taskId: %s",
             sessionId,
@@ -131,7 +131,7 @@
      * having session id [sessionId]
      */
     fun logTaskInfoChanged(sessionId: Int, taskUpdate: TaskUpdate) {
-        KtProtoLog.v(
+        ProtoLog.v(
             ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
             "DesktopModeLogger: Logging task info changed, session: %s taskId: %s",
             sessionId,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeLoggerTransitionObserver.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeLoggerTransitionObserver.kt
index 275f725d..066b5ad 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeLoggerTransitionObserver.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeLoggerTransitionObserver.kt
@@ -50,7 +50,6 @@
 import com.android.wm.shell.shared.TransitionUtil
 import com.android.wm.shell.sysui.ShellInit
 import com.android.wm.shell.transition.Transitions
-import com.android.wm.shell.util.KtProtoLog
 
 /**
  * A [Transitions.TransitionObserver] that observes transitions and the proposed changes to log
@@ -106,7 +105,7 @@
     ) {
         // this was a new recents animation
         if (info.isExitToRecentsTransition() && tasksSavedForRecents.isEmpty()) {
-            KtProtoLog.v(
+            ProtoLog.v(
                 WM_SHELL_DESKTOP_MODE,
                 "DesktopModeLogger: Recents animation running, saving tasks for later"
             )
@@ -132,7 +131,7 @@
                 info.flags == 0 &&
                 tasksSavedForRecents.isNotEmpty()
         ) {
-            KtProtoLog.v(
+            ProtoLog.v(
                 WM_SHELL_DESKTOP_MODE,
                 "DesktopModeLogger: Canceled recents animation, restoring tasks"
             )
@@ -202,7 +201,7 @@
             }
         }
 
-        KtProtoLog.v(
+        ProtoLog.v(
             WM_SHELL_DESKTOP_MODE,
             "DesktopModeLogger: taskInfo map after processing changes %s",
             postTransitionFreeformTasks.size()
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt
index df79b15..ca05864 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt
@@ -26,8 +26,8 @@
 import androidx.core.util.forEach
 import androidx.core.util.keyIterator
 import androidx.core.util.valueIterator
+import com.android.internal.protolog.ProtoLog
 import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE
-import com.android.wm.shell.util.KtProtoLog
 import java.io.PrintWriter
 import java.util.concurrent.Executor
 import java.util.function.Consumer
@@ -142,7 +142,7 @@
 
         val added = displayData.getOrCreate(displayId).activeTasks.add(taskId)
         if (added) {
-            KtProtoLog.d(
+            ProtoLog.d(
                 WM_SHELL_DESKTOP_MODE,
                 "DesktopTaskRepo: add active task=%d displayId=%d",
                 taskId,
@@ -167,7 +167,7 @@
             }
         }
         if (result) {
-            KtProtoLog.d(WM_SHELL_DESKTOP_MODE, "DesktopTaskRepo: remove active task=%d", taskId)
+            ProtoLog.d(WM_SHELL_DESKTOP_MODE, "DesktopTaskRepo: remove active task=%d", taskId)
         }
         return result
     }
@@ -180,7 +180,7 @@
     fun addClosingTask(displayId: Int, taskId: Int): Boolean {
         val added = displayData.getOrCreate(displayId).closingTasks.add(taskId)
         if (added) {
-            KtProtoLog.d(
+            ProtoLog.d(
                 WM_SHELL_DESKTOP_MODE,
                 "DesktopTaskRepo: added closing task=%d displayId=%d",
                 taskId,
@@ -203,7 +203,7 @@
             }
         }
         if (removed) {
-            KtProtoLog.d(WM_SHELL_DESKTOP_MODE, "DesktopTaskRepo: remove closing task=%d", taskId)
+            ProtoLog.d(WM_SHELL_DESKTOP_MODE, "DesktopTaskRepo: remove closing task=%d", taskId)
         }
         return removed
     }
@@ -316,14 +316,14 @@
 
         // Check if count changed
         if (prevCount != newCount) {
-            KtProtoLog.d(
+            ProtoLog.d(
                 WM_SHELL_DESKTOP_MODE,
                 "DesktopTaskRepo: update task visibility taskId=%d visible=%b displayId=%d",
                 taskId,
                 visible,
                 displayId
             )
-            KtProtoLog.d(
+            ProtoLog.d(
                 WM_SHELL_DESKTOP_MODE,
                 "DesktopTaskRepo: visibleTaskCount has changed from %d to %d",
                 prevCount,
@@ -341,7 +341,7 @@
 
     /** Get number of tasks that are marked as visible on given [displayId] */
     fun getVisibleTaskCount(displayId: Int): Int {
-        KtProtoLog.d(
+        ProtoLog.d(
             WM_SHELL_DESKTOP_MODE,
             "DesktopTaskRepo: visibleTaskCount= %d",
             displayData[displayId]?.visibleTasks?.size ?: 0
@@ -353,7 +353,7 @@
     // TODO(b/342417921): Identify if there is additional checks needed to move tasks for
     // multi-display scenarios.
     fun addOrMoveFreeformTaskToTop(displayId: Int, taskId: Int) {
-        KtProtoLog.d(
+        ProtoLog.d(
             WM_SHELL_DESKTOP_MODE,
             "DesktopTaskRepo: add or move task to top: display=%d, taskId=%d",
             displayId,
@@ -365,7 +365,7 @@
 
     /** Mark a Task as minimized. */
     fun minimizeTask(displayId: Int, taskId: Int) {
-        KtProtoLog.v(
+        ProtoLog.v(
             WM_SHELL_DESKTOP_MODE,
             "DesktopModeTaskRepository: minimize Task: display=%d, task=%d",
             displayId,
@@ -376,7 +376,7 @@
 
     /** Mark a Task as non-minimized. */
     fun unminimizeTask(displayId: Int, taskId: Int) {
-        KtProtoLog.v(
+        ProtoLog.v(
             WM_SHELL_DESKTOP_MODE,
             "DesktopModeTaskRepository: unminimize Task: display=%d, task=%d",
             displayId,
@@ -387,7 +387,7 @@
 
     /** Remove the task from the ordered list. */
     fun removeFreeformTask(displayId: Int, taskId: Int) {
-        KtProtoLog.d(
+        ProtoLog.d(
             WM_SHELL_DESKTOP_MODE,
             "DesktopTaskRepo: remove freeform task from ordered list: display=%d, taskId=%d",
             displayId,
@@ -395,7 +395,7 @@
         )
         displayData[displayId]?.freeformTasksInZOrder?.remove(taskId)
         boundsBeforeMaximizeByTaskId.remove(taskId)
-        KtProtoLog.d(
+        ProtoLog.d(
             WM_SHELL_DESKTOP_MODE,
             "DesktopTaskRepo: remaining freeform tasks: %s",
             displayData[displayId]?.freeformTasksInZOrder?.toDumpString() ?: ""
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
index 985901d..18157d6 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
@@ -50,6 +50,7 @@
 import androidx.annotation.BinderThread
 import com.android.internal.annotations.VisibleForTesting
 import com.android.internal.policy.ScreenDecorationsUtils
+import com.android.internal.protolog.ProtoLog
 import com.android.window.flags.Flags
 import com.android.wm.shell.RootTaskDisplayAreaOrganizer
 import com.android.wm.shell.ShellTaskOrganizer
@@ -88,7 +89,6 @@
 import com.android.wm.shell.sysui.ShellSharedConstants
 import com.android.wm.shell.transition.OneShotRemoteHandler
 import com.android.wm.shell.transition.Transitions
-import com.android.wm.shell.util.KtProtoLog
 import com.android.wm.shell.windowdecor.DragPositioningCallbackUtility
 import com.android.wm.shell.windowdecor.MoveToDesktopAnimator
 import com.android.wm.shell.windowdecor.OnTaskResizeAnimationListener
@@ -186,7 +186,7 @@
     }
 
     private fun onInit() {
-        KtProtoLog.d(WM_SHELL_DESKTOP_MODE, "Initialize DesktopTasksController")
+        ProtoLog.d(WM_SHELL_DESKTOP_MODE, "Initialize DesktopTasksController")
         shellCommandHandler.addDumpCallback(this::dump, this)
         shellCommandHandler.addCommandCallback("desktopmode", desktopModeShellCommandHandler, this)
         shellController.addExternalInterface(
@@ -200,7 +200,7 @@
         recentsTransitionHandler.addTransitionStateListener(
             object : RecentsTransitionStateListener {
                 override fun onAnimationStateChanged(running: Boolean) {
-                    KtProtoLog.v(
+                    ProtoLog.v(
                         WM_SHELL_DESKTOP_MODE,
                         "DesktopTasksController: recents animation state changed running=%b",
                         running
@@ -231,7 +231,7 @@
 
     /** Show all tasks, that are part of the desktop, on top of launcher */
     fun showDesktopApps(displayId: Int, remoteTransition: RemoteTransition? = null) {
-        KtProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: showDesktopApps")
+        ProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: showDesktopApps")
         val wct = WindowContainerTransaction()
         bringDesktopAppsToFront(displayId, wct)
 
@@ -282,7 +282,7 @@
                     moveToDesktop(allFocusedTasks[0].taskId, transitionSource = transitionSource)
                 }
                 else -> {
-                    KtProtoLog.w(
+                    ProtoLog.w(
                         WM_SHELL_DESKTOP_MODE,
                         "DesktopTasksController: Cannot enter desktop, expected less " +
                             "than 3 focused tasks but found %d",
@@ -312,7 +312,7 @@
         transitionSource: DesktopModeTransitionSource,
     ): Boolean {
         recentTasksController?.findTaskInBackground(taskId)?.let {
-            KtProtoLog.v(
+            ProtoLog.v(
                 WM_SHELL_DESKTOP_MODE,
                 "DesktopTasksController: moveToDesktopFromNonRunningTask taskId=%d",
                 taskId
@@ -346,14 +346,14 @@
     ) {
         if (Flags.enableDesktopWindowingModalsPolicy()
             && isTopActivityExemptFromDesktopWindowing(context, task)) {
-            KtProtoLog.w(
+            ProtoLog.w(
                 WM_SHELL_DESKTOP_MODE,
                 "DesktopTasksController: Cannot enter desktop, " +
                         "ineligible top activity found."
             )
             return
         }
-        KtProtoLog.v(
+        ProtoLog.v(
             WM_SHELL_DESKTOP_MODE,
             "DesktopTasksController: moveToDesktop taskId=%d",
             task.taskId
@@ -380,7 +380,7 @@
         taskInfo: RunningTaskInfo,
         dragToDesktopValueAnimator: MoveToDesktopAnimator,
     ) {
-        KtProtoLog.v(
+        ProtoLog.v(
             WM_SHELL_DESKTOP_MODE,
             "DesktopTasksController: startDragToDesktop taskId=%d",
             taskInfo.taskId
@@ -396,7 +396,7 @@
      * [startDragToDesktop].
      */
     private fun finalizeDragToDesktop(taskInfo: RunningTaskInfo, freeformBounds: Rect) {
-        KtProtoLog.v(
+        ProtoLog.v(
             WM_SHELL_DESKTOP_MODE,
             "DesktopTasksController: finalizeDragToDesktop taskId=%d",
             taskInfo.taskId
@@ -440,7 +440,7 @@
         }
         if (!desktopModeTaskRepository.addClosingTask(displayId, taskId)) {
             // Could happen if the task hasn't been removed from closing list after it disappeared
-            KtProtoLog.w(
+            ProtoLog.w(
                 WM_SHELL_DESKTOP_MODE,
                 "DesktopTasksController: the task with taskId=%d is already closing!",
                 taskId
@@ -464,7 +464,7 @@
 
     /** Move a desktop app to split screen. */
     fun moveToSplit(task: RunningTaskInfo) {
-        KtProtoLog.v(
+        ProtoLog.v(
             WM_SHELL_DESKTOP_MODE,
             "DesktopTasksController: moveToSplit taskId=%d",
             task.taskId
@@ -497,7 +497,7 @@
      * [startDragToDesktop].
      */
     fun cancelDragToDesktop(task: RunningTaskInfo) {
-        KtProtoLog.v(
+        ProtoLog.v(
             WM_SHELL_DESKTOP_MODE,
             "DesktopTasksController: cancelDragToDesktop taskId=%d",
             task.taskId
@@ -512,7 +512,7 @@
         position: Point,
         transitionSource: DesktopModeTransitionSource
     ) {
-        KtProtoLog.v(
+        ProtoLog.v(
             WM_SHELL_DESKTOP_MODE,
             "DesktopTasksController: moveToFullscreen with animation taskId=%d",
             task.taskId
@@ -540,7 +540,7 @@
 
     /** Move a task to the front */
     fun moveTaskToFront(taskInfo: RunningTaskInfo) {
-        KtProtoLog.v(
+        ProtoLog.v(
             WM_SHELL_DESKTOP_MODE,
             "DesktopTasksController: moveTaskToFront taskId=%d",
             taskInfo.taskId
@@ -571,10 +571,10 @@
     fun moveToNextDisplay(taskId: Int) {
         val task = shellTaskOrganizer.getRunningTaskInfo(taskId)
         if (task == null) {
-            KtProtoLog.w(WM_SHELL_DESKTOP_MODE, "moveToNextDisplay: taskId=%d not found", taskId)
+            ProtoLog.w(WM_SHELL_DESKTOP_MODE, "moveToNextDisplay: taskId=%d not found", taskId)
             return
         }
-        KtProtoLog.v(
+        ProtoLog.v(
             WM_SHELL_DESKTOP_MODE,
             "moveToNextDisplay: taskId=%d taskDisplayId=%d",
             taskId,
@@ -589,7 +589,7 @@
             newDisplayId = displayIds.firstOrNull { displayId -> displayId < task.displayId }
         }
         if (newDisplayId == null) {
-            KtProtoLog.w(WM_SHELL_DESKTOP_MODE, "moveToNextDisplay: next display not found")
+            ProtoLog.w(WM_SHELL_DESKTOP_MODE, "moveToNextDisplay: next display not found")
             return
         }
         moveToDisplay(task, newDisplayId)
@@ -601,7 +601,7 @@
      * No-op if task is already on that display per [RunningTaskInfo.displayId].
      */
     private fun moveToDisplay(task: RunningTaskInfo, displayId: Int) {
-        KtProtoLog.v(
+        ProtoLog.v(
             WM_SHELL_DESKTOP_MODE,
             "moveToDisplay: taskId=%d displayId=%d",
             task.taskId,
@@ -609,13 +609,13 @@
         )
 
         if (task.displayId == displayId) {
-            KtProtoLog.d(WM_SHELL_DESKTOP_MODE, "moveToDisplay: task already on display")
+            ProtoLog.d(WM_SHELL_DESKTOP_MODE, "moveToDisplay: task already on display")
             return
         }
 
         val displayAreaInfo = rootTaskDisplayAreaOrganizer.getDisplayAreaInfo(displayId)
         if (displayAreaInfo == null) {
-            KtProtoLog.w(WM_SHELL_DESKTOP_MODE, "moveToDisplay: display not found")
+            ProtoLog.w(WM_SHELL_DESKTOP_MODE, "moveToDisplay: display not found")
             return
         }
 
@@ -770,7 +770,7 @@
         wct: WindowContainerTransaction,
         newTaskIdInFront: Int? = null
     ): RunningTaskInfo? {
-        KtProtoLog.v(
+        ProtoLog.v(
             WM_SHELL_DESKTOP_MODE,
             "DesktopTasksController: bringDesktopAppsToFront, newTaskIdInFront=%s",
             newTaskIdInFront ?: "null"
@@ -815,7 +815,7 @@
     }
 
     private fun addWallpaperActivity(wct: WindowContainerTransaction) {
-        KtProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: addWallpaper")
+        ProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: addWallpaper")
         val intent = Intent(context, DesktopWallpaperActivity::class.java)
         val options =
             ActivityOptions.makeBasic().apply {
@@ -835,7 +835,7 @@
 
     private fun removeWallpaperActivity(wct: WindowContainerTransaction) {
         desktopModeTaskRepository.wallpaperActivityToken?.let { token ->
-            KtProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: removeWallpaper")
+            ProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: removeWallpaper")
             wct.removeTask(token)
         }
     }
@@ -873,7 +873,7 @@
         transition: IBinder,
         request: TransitionRequestInfo
     ): WindowContainerTransaction? {
-        KtProtoLog.v(
+        ProtoLog.v(
             WM_SHELL_DESKTOP_MODE,
             "DesktopTasksController: handleRequest request=%s",
             request
@@ -915,7 +915,7 @@
             }
 
         if (!shouldHandleRequest) {
-            KtProtoLog.v(
+            ProtoLog.v(
                 WM_SHELL_DESKTOP_MODE,
                 "DesktopTasksController: skipping handleRequest reason=%s",
                 reason
@@ -939,7 +939,7 @@
                     }
                 }
             }
-        KtProtoLog.v(
+        ProtoLog.v(
             WM_SHELL_DESKTOP_MODE,
             "DesktopTasksController: handleRequest result=%s",
             result ?: "null"
@@ -977,15 +977,15 @@
         task: RunningTaskInfo,
         transition: IBinder
     ): WindowContainerTransaction? {
-        KtProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: handleFreeformTaskLaunch")
+        ProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: handleFreeformTaskLaunch")
         if (keyguardManager.isKeyguardLocked) {
             // Do NOT handle freeform task launch when locked.
             // It will be launched in fullscreen windowing mode (Details: b/160925539)
-            KtProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: skip keyguard is locked")
+            ProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: skip keyguard is locked")
             return null
         }
         if (!desktopModeTaskRepository.isDesktopModeShowing(task.displayId)) {
-            KtProtoLog.d(
+            ProtoLog.d(
                 WM_SHELL_DESKTOP_MODE,
                 "DesktopTasksController: bring desktop tasks to front on transition" +
                     " taskId=%d",
@@ -1014,9 +1014,9 @@
         task: RunningTaskInfo,
         transition: IBinder
     ): WindowContainerTransaction? {
-        KtProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: handleFullscreenTaskLaunch")
+        ProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: handleFullscreenTaskLaunch")
         if (desktopModeTaskRepository.isDesktopModeShowing(task.displayId)) {
-            KtProtoLog.d(
+            ProtoLog.d(
                 WM_SHELL_DESKTOP_MODE,
                 "DesktopTasksController: switch fullscreen task to freeform on transition" +
                     " taskId=%d",
@@ -1059,7 +1059,7 @@
         }
         if (!desktopModeTaskRepository.addClosingTask(task.displayId, task.taskId)) {
             // Could happen if the task hasn't been removed from closing list after it disappeared
-            KtProtoLog.w(
+            ProtoLog.w(
                 WM_SHELL_DESKTOP_MODE,
                 "DesktopTasksController: the task with taskId=%d is already closing!",
                 task.taskId
@@ -1398,7 +1398,7 @@
         if (!multiInstanceHelper.supportsMultiInstanceSplit(launchComponent)) {
             // TODO(b/320797628): Should only return early if there is an existing running task, and
             //                    notify the user as well. But for now, just ignore the drop.
-            KtProtoLog.v(WM_SHELL_DESKTOP_MODE, "Dropped intent does not support multi-instance")
+            ProtoLog.v(WM_SHELL_DESKTOP_MODE, "Dropped intent does not support multi-instance")
             return false
         }
 
@@ -1489,7 +1489,7 @@
         private val listener: VisibleTasksListener =
             object : VisibleTasksListener {
                 override fun onTasksVisibilityChanged(displayId: Int, visibleTasksCount: Int) {
-                    KtProtoLog.v(
+                    ProtoLog.v(
                         WM_SHELL_DESKTOP_MODE,
                         "IDesktopModeImpl: onVisibilityChanged display=%d visible=%d",
                         displayId,
@@ -1534,11 +1534,11 @@
         }
 
         override fun stashDesktopApps(displayId: Int) {
-            KtProtoLog.w(WM_SHELL_DESKTOP_MODE, "IDesktopModeImpl: stashDesktopApps is deprecated")
+            ProtoLog.w(WM_SHELL_DESKTOP_MODE, "IDesktopModeImpl: stashDesktopApps is deprecated")
         }
 
         override fun hideStashedDesktopApps(displayId: Int) {
-            KtProtoLog.w(
+            ProtoLog.w(
                 WM_SHELL_DESKTOP_MODE,
                 "IDesktopModeImpl: hideStashedDesktopApps is deprecated"
             )
@@ -1565,7 +1565,7 @@
         }
 
         override fun setTaskListener(listener: IDesktopTaskListener?) {
-            KtProtoLog.v(
+            ProtoLog.v(
                 WM_SHELL_DESKTOP_MODE,
                 "IDesktopModeImpl: set task listener=%s",
                 listener ?: "null"
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksLimiter.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksLimiter.kt
index 3f5bd1a..534cc22 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksLimiter.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksLimiter.kt
@@ -23,12 +23,12 @@
 import android.window.TransitionInfo
 import android.window.WindowContainerTransaction
 import androidx.annotation.VisibleForTesting
+import com.android.internal.protolog.ProtoLog
 import com.android.wm.shell.ShellTaskOrganizer
 import com.android.wm.shell.protolog.ShellProtoLogGroup
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import com.android.wm.shell.transition.Transitions
 import com.android.wm.shell.transition.Transitions.TransitionObserver
-import com.android.wm.shell.util.KtProtoLog
 
 /**
  * Limits the number of tasks shown in Desktop Mode.
@@ -71,7 +71,7 @@
             if (!taskRepository.isActiveTask(taskToMinimize.taskId)) return
 
             if (!isTaskReorderedToBackOrInvisible(info, taskToMinimize)) {
-                KtProtoLog.v(
+                ProtoLog.v(
                         ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
                         "DesktopTasksLimiter: task %d is not reordered to back nor invis",
                         taskToMinimize.taskId)
@@ -109,7 +109,7 @@
         }
 
         override fun onTransitionFinished(transition: IBinder, aborted: Boolean) {
-            KtProtoLog.v(
+            ProtoLog.v(
                     ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
                     "DesktopTasksLimiter: transition %s finished", transition)
             mPendingTransitionTokensAndTasks.remove(transition)
@@ -133,7 +133,7 @@
             if (remainingMinimizedTasks.isEmpty()) {
                 return
             }
-            KtProtoLog.v(
+            ProtoLog.v(
                 ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
                 "DesktopTasksLimiter: removing leftover minimized tasks: $remainingMinimizedTasks")
             remainingMinimizedTasks.forEach { taskIdToRemove ->
@@ -150,7 +150,7 @@
      * finished so we don't minimize the task if the transition fails.
      */
     private fun markTaskMinimized(displayId: Int, taskId: Int) {
-        KtProtoLog.v(
+        ProtoLog.v(
                 ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
                 "DesktopTasksLimiter: marking %d as minimized", taskId)
         taskRepository.minimizeTask(displayId, taskId)
@@ -169,7 +169,7 @@
             wct: WindowContainerTransaction,
             newFrontTaskInfo: RunningTaskInfo,
     ): RunningTaskInfo? {
-        KtProtoLog.v(
+        ProtoLog.v(
                 ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
                 "DesktopTasksLimiter: addMinimizeBackTaskChangesIfNeeded, newFrontTask=%d",
                 newFrontTaskInfo.taskId)
@@ -217,7 +217,7 @@
             visibleFreeformTaskIdsOrderedFrontToBack: List<Int>
     ): RunningTaskInfo? {
         if (visibleFreeformTaskIdsOrderedFrontToBack.size <= getMaxTaskLimit()) {
-            KtProtoLog.v(
+            ProtoLog.v(
                     ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
                     "DesktopTasksLimiter: no need to minimize; tasks below limit")
             // No need to minimize anything
@@ -227,7 +227,7 @@
                 shellTaskOrganizer.getRunningTaskInfo(
                         visibleFreeformTaskIdsOrderedFrontToBack.last())
         if (taskToMinimize == null) {
-            KtProtoLog.e(
+            ProtoLog.e(
                     ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
                     "DesktopTasksLimiter: taskToMinimize == null")
             return null
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserver.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserver.kt
index f01f645..45ed7565 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserver.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserver.kt
@@ -21,12 +21,12 @@
 import android.view.SurfaceControl
 import android.view.WindowManager
 import android.window.TransitionInfo
+import com.android.internal.protolog.ProtoLog
 import com.android.window.flags.Flags.enableDesktopWindowingWallpaperActivity
 import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import com.android.wm.shell.sysui.ShellInit
 import com.android.wm.shell.transition.Transitions
-import com.android.wm.shell.util.KtProtoLog
 
 /**
  * A [Transitions.TransitionObserver] that observes shell transitions and updates
@@ -49,7 +49,7 @@
     }
 
     fun onInit() {
-        KtProtoLog.d(WM_SHELL_DESKTOP_MODE, "DesktopTasksTransitionObserver: onInit")
+        ProtoLog.d(WM_SHELL_DESKTOP_MODE, "DesktopTasksTransitionObserver: onInit")
         transitions.registerObserver(this)
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopWallpaperActivity.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopWallpaperActivity.kt
index c4a4474..1c2415c 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopWallpaperActivity.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopWallpaperActivity.kt
@@ -21,8 +21,8 @@
 import android.content.ComponentName
 import android.os.Bundle
 import android.view.WindowManager
+import com.android.internal.protolog.ProtoLog
 import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE
-import com.android.wm.shell.util.KtProtoLog
 
 /**
  * A transparent activity used in the desktop mode to show the wallpaper under the freeform windows.
@@ -36,7 +36,7 @@
 class DesktopWallpaperActivity : Activity() {
 
     override fun onCreate(savedInstanceState: Bundle?) {
-        KtProtoLog.d(WM_SHELL_DESKTOP_MODE, "DesktopWallpaperActivity: onCreate")
+        ProtoLog.d(WM_SHELL_DESKTOP_MODE, "DesktopWallpaperActivity: onCreate")
         super.onCreate(savedInstanceState)
         window.addFlags(WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE)
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandler.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandler.kt
index d99b724..ddee8fa 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandler.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandler.kt
@@ -29,6 +29,7 @@
 import android.window.TransitionRequestInfo
 import android.window.WindowContainerToken
 import android.window.WindowContainerTransaction
+import com.android.internal.protolog.ProtoLog
 import com.android.wm.shell.RootTaskDisplayAreaOrganizer
 import com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_BOTTOM_OR_RIGHT
 import com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_TOP_OR_LEFT
@@ -42,7 +43,6 @@
 import com.android.wm.shell.transition.Transitions.TRANSIT_DESKTOP_MODE_END_DRAG_TO_DESKTOP
 import com.android.wm.shell.transition.Transitions.TRANSIT_DESKTOP_MODE_START_DRAG_TO_DESKTOP
 import com.android.wm.shell.transition.Transitions.TransitionHandler
-import com.android.wm.shell.util.KtProtoLog
 import com.android.wm.shell.windowdecor.MoveToDesktopAnimator
 import com.android.wm.shell.windowdecor.MoveToDesktopAnimator.Companion.DRAG_FREEFORM_SCALE
 import com.android.wm.shell.windowdecor.OnTaskResizeAnimationListener
@@ -114,7 +114,7 @@
         dragToDesktopAnimator: MoveToDesktopAnimator,
     ) {
         if (inProgress) {
-            KtProtoLog.v(
+            ProtoLog.v(
                 ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
                 "DragToDesktop: Drag to desktop transition already in progress."
             )
@@ -599,7 +599,7 @@
     ) {
         val state = transitionState ?: return
         if (aborted && state.startTransitionToken == transition) {
-            KtProtoLog.v(
+            ProtoLog.v(
                 ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
                 "DragToDesktop: onTransitionConsumed() start transition aborted"
             )
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/sysui/ShellInit.java b/libs/WindowManager/Shell/src/com/android/wm/shell/sysui/ShellInit.java
index 3a68009..dd4595a 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/sysui/ShellInit.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/sysui/ShellInit.java
@@ -27,6 +27,7 @@
 
 import com.android.internal.protolog.ProtoLog;
 import com.android.wm.shell.common.ShellExecutor;
+import com.android.wm.shell.protolog.ShellProtoLogGroup;
 
 import java.util.ArrayList;
 
@@ -75,6 +76,7 @@
      */
     @VisibleForTesting
     public void init() {
+        ProtoLog.registerGroups(ShellProtoLogGroup.values());
         ProtoLog.v(WM_SHELL_INIT, "Initializing Shell Components: %d", mInitCallbacks.size());
         SurfaceControl.setDebugUsageAfterRelease(true);
         // Init in order of registration
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/util/KtProtoLog.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/util/KtProtoLog.kt
deleted file mode 100644
index ee6c81a..0000000
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/util/KtProtoLog.kt
+++ /dev/null
@@ -1,74 +0,0 @@
-/*
- * Copyright (C) 2023 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.wm.shell.util
-
-import android.util.Log
-import com.android.internal.protolog.common.IProtoLogGroup
-import com.android.internal.protolog.ProtoLog
-
-/**
- * Log messages using an API similar to [com.android.internal.protolog.ProtoLog]. Useful for
- * logging from Kotlin classes as ProtoLog does not have support for Kotlin.
- *
- * All messages are logged to logcat if logging is enabled for that [IProtoLogGroup].
- */
-// TODO(b/168581922): remove once ProtoLog adds support for Kotlin
-class KtProtoLog {
-    companion object {
-        /** @see [com.android.internal.protolog.ProtoLog.d] */
-        fun d(group: IProtoLogGroup, messageString: String, vararg args: Any) {
-            if (group.isLogToLogcat) {
-                Log.d(group.tag, String.format(messageString, *args))
-            }
-        }
-
-        /** @see [com.android.internal.protolog.ProtoLog.v] */
-        fun v(group: IProtoLogGroup, messageString: String, vararg args: Any) {
-            if (group.isLogToLogcat) {
-                Log.v(group.tag, String.format(messageString, *args))
-            }
-        }
-
-        /** @see [com.android.internal.protolog.ProtoLog.i] */
-        fun i(group: IProtoLogGroup, messageString: String, vararg args: Any) {
-            if (group.isLogToLogcat) {
-                Log.i(group.tag, String.format(messageString, *args))
-            }
-        }
-
-        /** @see [com.android.internal.protolog.ProtoLog.w] */
-        fun w(group: IProtoLogGroup, messageString: String, vararg args: Any) {
-            if (group.isLogToLogcat) {
-                Log.w(group.tag, String.format(messageString, *args))
-            }
-        }
-
-        /** @see [com.android.internal.protolog.ProtoLog.e] */
-        fun e(group: IProtoLogGroup, messageString: String, vararg args: Any) {
-            if (group.isLogToLogcat) {
-                Log.e(group.tag, String.format(messageString, *args))
-            }
-        }
-
-        /** @see [com.android.internal.protolog.ProtoLog.wtf] */
-        fun wtf(group: IProtoLogGroup, messageString: String, vararg args: Any) {
-            if (group.isLogToLogcat) {
-                Log.wtf(group.tag, String.format(messageString, *args))
-            }
-        }
-    }
-}
diff --git a/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/flicker/DesktopModeFlickerScenarios.kt b/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/flicker/DesktopModeFlickerScenarios.kt
index c725b08..430f80b 100644
--- a/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/flicker/DesktopModeFlickerScenarios.kt
+++ b/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/flicker/DesktopModeFlickerScenarios.kt
@@ -20,6 +20,7 @@
 import android.tools.flicker.assertors.assertions.AppLayerIsInvisibleAtEnd
 import android.tools.flicker.assertors.assertions.AppLayerIsVisibleAlways
 import android.tools.flicker.assertors.assertions.AppLayerIsVisibleAtStart
+import android.tools.flicker.assertors.assertions.AppWindowBecomesVisible
 import android.tools.flicker.assertors.assertions.AppWindowHasDesktopModeInitialBoundsAtTheEnd
 import android.tools.flicker.assertors.assertions.AppWindowHasSizeOfAtLeast
 import android.tools.flicker.assertors.assertions.AppWindowIsInvisibleAtEnd
@@ -45,27 +46,30 @@
             FlickerConfigEntry(
                 scenarioId = ScenarioId("END_DRAG_TO_DESKTOP"),
                 extractor =
-                ShellTransitionScenarioExtractor(
-                    transitionMatcher =
-                    object : ITransitionMatcher {
-                        override fun findAll(
-                            transitions: Collection<Transition>
-                        ): Collection<Transition> {
-                            return transitions.filter {
-                                it.type == TransitionType.DESKTOP_MODE_END_DRAG_TO_DESKTOP
+                    ShellTransitionScenarioExtractor(
+                        transitionMatcher =
+                            object : ITransitionMatcher {
+                                override fun findAll(
+                                    transitions: Collection<Transition>
+                                ): Collection<Transition> {
+                                    return transitions.filter {
+                                        // TODO(351168217) Use jank CUJ to extract a longer trace
+                                        it.type == TransitionType.DESKTOP_MODE_END_DRAG_TO_DESKTOP
+                                    }
+                                }
                             }
-                        }
-                    }
-                ),
+                    ),
                 assertions =
-                AssertionTemplates.COMMON_ASSERTIONS +
+                    AssertionTemplates.COMMON_ASSERTIONS +
                         listOf(
-                            AppLayerIsVisibleAlways(Components.DESKTOP_MODE_APP),
-                            AppWindowOnTopAtEnd(Components.DESKTOP_MODE_APP),
-                            AppWindowHasDesktopModeInitialBoundsAtTheEnd(
-                                Components.DESKTOP_MODE_APP
+                                AppLayerIsVisibleAlways(Components.DESKTOP_MODE_APP),
+                                AppWindowOnTopAtEnd(Components.DESKTOP_MODE_APP),
+                                AppWindowHasDesktopModeInitialBoundsAtTheEnd(
+                                    Components.DESKTOP_MODE_APP
+                                ),
+                                AppWindowBecomesVisible(DESKTOP_WALLPAPER)
                             )
-                        ).associateBy({ it }, { AssertionInvocationGroup.BLOCKING }),
+                            .associateBy({ it }, { AssertionInvocationGroup.BLOCKING }),
             )
 
         // Use this scenario for closing an app in desktop windowing, except the last app. For the
@@ -74,63 +78,65 @@
             FlickerConfigEntry(
                 scenarioId = ScenarioId("CLOSE_APP"),
                 extractor =
-                ShellTransitionScenarioExtractor(
-                    transitionMatcher =
-                    object : ITransitionMatcher {
-                        override fun findAll(
-                            transitions: Collection<Transition>
-                        ): Collection<Transition> {
-                            // In case there are multiple windows closing, filter out the
-                            // last window closing. It should use the CLOSE_LAST_APP
-                            // scenario below.
-                            return transitions
-                                    .filter { it.type == TransitionType.CLOSE }
-                                    .sortedByDescending { it.id }
-                                    .drop(1)
-                        }
-                    }
-                ),
+                    ShellTransitionScenarioExtractor(
+                        transitionMatcher =
+                            object : ITransitionMatcher {
+                                override fun findAll(
+                                    transitions: Collection<Transition>
+                                ): Collection<Transition> {
+                                    // In case there are multiple windows closing, filter out the
+                                    // last window closing. It should use the CLOSE_LAST_APP
+                                    // scenario below.
+                                    return transitions
+                                        .filter { it.type == TransitionType.CLOSE }
+                                        .sortedByDescending { it.id }
+                                        .drop(1)
+                                }
+                            }
+                    ),
                 assertions =
-                AssertionTemplates.COMMON_ASSERTIONS +
+                    AssertionTemplates.COMMON_ASSERTIONS +
                         listOf(
-                            AppWindowOnTopAtStart(Components.DESKTOP_MODE_APP),
-                            AppLayerIsVisibleAtStart(Components.DESKTOP_MODE_APP),
-                            AppLayerIsInvisibleAtEnd(Components.DESKTOP_MODE_APP),
-                        ).associateBy({ it }, { AssertionInvocationGroup.BLOCKING }),
+                                AppWindowOnTopAtStart(Components.DESKTOP_MODE_APP),
+                                AppLayerIsVisibleAtStart(Components.DESKTOP_MODE_APP),
+                                AppLayerIsInvisibleAtEnd(Components.DESKTOP_MODE_APP),
+                            )
+                            .associateBy({ it }, { AssertionInvocationGroup.BLOCKING }),
             )
 
         val CLOSE_LAST_APP =
             FlickerConfigEntry(
                 scenarioId = ScenarioId("CLOSE_LAST_APP"),
                 extractor =
-                ShellTransitionScenarioExtractor(
-                    transitionMatcher =
-                    object : ITransitionMatcher {
-                        override fun findAll(
-                            transitions: Collection<Transition>
-                        ): Collection<Transition> {
-                            val lastTransition =
-                                transitions
-                                        .filter { it.type == TransitionType.CLOSE }
-                                        .maxByOrNull { it.id }!!
-                            return listOf(lastTransition)
-                        }
-                    }
-                ),
+                    ShellTransitionScenarioExtractor(
+                        transitionMatcher =
+                            object : ITransitionMatcher {
+                                override fun findAll(
+                                    transitions: Collection<Transition>
+                                ): Collection<Transition> {
+                                    val lastTransition =
+                                        transitions
+                                            .filter { it.type == TransitionType.CLOSE }
+                                            .maxByOrNull { it.id }!!
+                                    return listOf(lastTransition)
+                                }
+                            }
+                    ),
                 assertions =
-                AssertionTemplates.COMMON_ASSERTIONS +
+                    AssertionTemplates.COMMON_ASSERTIONS +
                         listOf(
-                            AppWindowIsInvisibleAtEnd(Components.DESKTOP_MODE_APP),
-                            LauncherWindowReplacesAppAsTopWindow(Components.DESKTOP_MODE_APP),
-                            AppWindowIsInvisibleAtEnd(DESKTOP_WALLPAPER)
-                        ).associateBy({ it }, { AssertionInvocationGroup.BLOCKING }),
+                                AppWindowIsInvisibleAtEnd(Components.DESKTOP_MODE_APP),
+                                LauncherWindowReplacesAppAsTopWindow(Components.DESKTOP_MODE_APP),
+                                AppWindowIsInvisibleAtEnd(DESKTOP_WALLPAPER)
+                            )
+                            .associateBy({ it }, { AssertionInvocationGroup.BLOCKING }),
             )
 
         val CORNER_RESIZE =
             FlickerConfigEntry(
                 scenarioId = ScenarioId("CORNER_RESIZE"),
                 extractor =
-                TaggedScenarioExtractorBuilder()
+                    TaggedScenarioExtractorBuilder()
                         .setTargetTag(CujType.CUJ_DESKTOP_MODE_RESIZE_WINDOW)
                         .setTransitionMatcher(
                             TaggedCujTransitionMatcher(associatedTransitionRequired = false)
@@ -143,16 +149,16 @@
             FlickerConfigEntry(
                 scenarioId = ScenarioId("CORNER_RESIZE_TO_MINIMUM_SIZE"),
                 extractor =
-                TaggedScenarioExtractorBuilder()
+                    TaggedScenarioExtractorBuilder()
                         .setTargetTag(CujType.CUJ_DESKTOP_MODE_RESIZE_WINDOW)
                         .setTransitionMatcher(
                             TaggedCujTransitionMatcher(associatedTransitionRequired = false)
-                        ).build(),
+                        )
+                        .build(),
                 assertions =
-                AssertionTemplates.DESKTOP_MODE_APP_VISIBILITY_ASSERTIONS +
-                        listOf(
-                            AppWindowHasSizeOfAtLeast(Components.DESKTOP_MODE_APP, 770, 700)
-                        ).associateBy({ it }, { AssertionInvocationGroup.BLOCKING }),
+                    AssertionTemplates.DESKTOP_MODE_APP_VISIBILITY_ASSERTIONS +
+                        listOf(AppWindowHasSizeOfAtLeast(Components.DESKTOP_MODE_APP, 770, 700))
+                            .associateBy({ it }, { AssertionInvocationGroup.BLOCKING }),
             )
     }
-}
\ No newline at end of file
+}
diff --git a/packages/SystemUI/Android.bp b/packages/SystemUI/Android.bp
index 469b9ce..9cbb1bd 100644
--- a/packages/SystemUI/Android.bp
+++ b/packages/SystemUI/Android.bp
@@ -555,6 +555,7 @@
         "androidx.exifinterface_exifinterface",
         "androidx.room_room-runtime",
         "androidx.room_room-ktx",
+        "androidx.datastore_datastore-preferences",
         "com.google.android.material_material",
         "device_state_flags_lib",
         "kotlinx_coroutines_android",
@@ -708,6 +709,7 @@
         "androidx.exifinterface_exifinterface",
         "androidx.room_room-runtime",
         "androidx.room_room-ktx",
+        "androidx.datastore_datastore-preferences",
         "device_state_flags_lib",
         "kotlinx-coroutines-android",
         "kotlinx-coroutines-core",
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/education/data/repository/ContextualEducationRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/education/data/repository/ContextualEducationRepositoryTest.kt
new file mode 100644
index 0000000..4a5342a
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/education/data/repository/ContextualEducationRepositoryTest.kt
@@ -0,0 +1,95 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.education.data.repository
+
+import android.content.Context
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.SysuiTestableContext
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.shared.education.GestureType.BACK_GESTURE
+import com.google.common.truth.Truth.assertThat
+import java.io.File
+import javax.inject.Provider
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.rules.TemporaryFolder
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class ContextualEducationRepositoryTest : SysuiTestCase() {
+
+    private lateinit var underTest: ContextualEducationRepository
+    private val kosmos = Kosmos()
+    private val testScope = kosmos.testScope
+    private val dsScopeProvider: Provider<CoroutineScope> = Provider {
+        TestScope(kosmos.testDispatcher).backgroundScope
+    }
+    private val testUserId = 1111
+
+    // For deleting any test files created after the test
+    @get:Rule val tmpFolder: TemporaryFolder = TemporaryFolder.builder().assureDeletion().build()
+
+    @Before
+    fun setUp() {
+        // Create TestContext here because TemporaryFolder.create() is called in @Before. It is
+        // needed before calling TemporaryFolder.newFolder().
+        val testContext = TestContext(context, tmpFolder.newFolder())
+        val userRepository = UserContextualEducationRepository(testContext, dsScopeProvider)
+        underTest = ContextualEducationRepository(userRepository)
+        underTest.setUser(testUserId)
+    }
+
+    @Test
+    fun changeRetrievedValueForNewUser() =
+        testScope.runTest {
+            // Update data for old user.
+            underTest.incrementSignalCount(BACK_GESTURE)
+            val model by collectLastValue(underTest.readGestureEduModelFlow(BACK_GESTURE))
+            assertThat(model?.signalCount).isEqualTo(1)
+
+            // User is changed.
+            underTest.setUser(1112)
+            // Assert count is 0 after user is changed.
+            assertThat(model?.signalCount).isEqualTo(0)
+        }
+
+    @Test
+    fun incrementSignalCount() =
+        testScope.runTest {
+            underTest.incrementSignalCount(BACK_GESTURE)
+            val model by collectLastValue(underTest.readGestureEduModelFlow(BACK_GESTURE))
+            assertThat(model?.signalCount).isEqualTo(1)
+        }
+
+    /** Test context which allows overriding getFilesDir path */
+    private class TestContext(context: Context, private val folder: File) :
+        SysuiTestableContext(context) {
+        override fun getFilesDir(): File {
+            return folder
+        }
+    }
+}
diff --git a/packages/SystemUI/shared/src/com/android/systemui/shared/education/GestureType.kt b/packages/SystemUI/shared/src/com/android/systemui/shared/education/GestureType.kt
new file mode 100644
index 0000000..9a5c77a
--- /dev/null
+++ b/packages/SystemUI/shared/src/com/android/systemui/shared/education/GestureType.kt
@@ -0,0 +1,21 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.shared.education
+
+enum class GestureType {
+    BACK_GESTURE,
+}
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java b/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
index 572283a..08cfd37 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
+++ b/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
@@ -61,6 +61,7 @@
 import com.android.systemui.doze.dagger.DozeComponent;
 import com.android.systemui.dreams.dagger.DreamModule;
 import com.android.systemui.dump.DumpManager;
+import com.android.systemui.education.dagger.ContextualEducationModule;
 import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.flags.FlagDependenciesModule;
 import com.android.systemui.flags.FlagsModule;
@@ -259,7 +260,8 @@
         UserModule.class,
         UtilModule.class,
         NoteTaskModule.class,
-        WalletModule.class
+        WalletModule.class,
+        ContextualEducationModule.class
         },
         subcomponents = {
             ComplicationComponent.class,
diff --git a/packages/SystemUI/src/com/android/systemui/education/dagger/ContextualEducationModule.kt b/packages/SystemUI/src/com/android/systemui/education/dagger/ContextualEducationModule.kt
new file mode 100644
index 0000000..e2bcb6b
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/education/dagger/ContextualEducationModule.kt
@@ -0,0 +1,40 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.education.dagger
+
+import com.android.systemui.dagger.qualifiers.Background
+import dagger.Module
+import dagger.Provides
+import javax.inject.Qualifier
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.SupervisorJob
+
+@Module
+interface ContextualEducationModule {
+    @Qualifier annotation class EduDataStoreScope
+
+    companion object {
+        @EduDataStoreScope
+        @Provides
+        fun provideEduDataStoreScope(
+            @Background bgDispatcher: CoroutineDispatcher
+        ): CoroutineScope {
+            return CoroutineScope(bgDispatcher + SupervisorJob())
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/education/data/model/GestureEduModel.kt b/packages/SystemUI/src/com/android/systemui/education/data/model/GestureEduModel.kt
new file mode 100644
index 0000000..af35e8c
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/education/data/model/GestureEduModel.kt
@@ -0,0 +1,26 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.education.data.model
+
+/**
+ * Model to store education data related to each gesture (e.g. Back, Home, All Apps, Overview). Each
+ * gesture stores its own model separately.
+ */
+data class GestureEduModel(
+    val signalCount: Int,
+    val educationShownCount: Int,
+)
diff --git a/packages/SystemUI/src/com/android/systemui/education/data/repository/ContextualEducationRepository.kt b/packages/SystemUI/src/com/android/systemui/education/data/repository/ContextualEducationRepository.kt
new file mode 100644
index 0000000..c9dd833
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/education/data/repository/ContextualEducationRepository.kt
@@ -0,0 +1,42 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.education.data.repository
+
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.shared.education.GestureType
+import javax.inject.Inject
+
+/**
+ * Provide methods to read and update on field level and allow setting datastore when user is
+ * changed
+ */
+@SysUISingleton
+class ContextualEducationRepository
+@Inject
+constructor(private val userEduRepository: UserContextualEducationRepository) {
+    /** To change data store when user is changed */
+    fun setUser(userId: Int) = userEduRepository.setUser(userId)
+
+    fun readGestureEduModelFlow(gestureType: GestureType) =
+        userEduRepository.readGestureEduModelFlow(gestureType)
+
+    suspend fun incrementSignalCount(gestureType: GestureType) {
+        userEduRepository.updateGestureEduModel(gestureType) {
+            it.copy(signalCount = it.signalCount + 1)
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/education/data/repository/UserContextualEducationRepository.kt b/packages/SystemUI/src/com/android/systemui/education/data/repository/UserContextualEducationRepository.kt
new file mode 100644
index 0000000..229511a
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/education/data/repository/UserContextualEducationRepository.kt
@@ -0,0 +1,114 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.education.data.repository
+
+import android.content.Context
+import androidx.datastore.core.DataStore
+import androidx.datastore.preferences.core.PreferenceDataStoreFactory
+import androidx.datastore.preferences.core.Preferences
+import androidx.datastore.preferences.core.edit
+import androidx.datastore.preferences.core.intPreferencesKey
+import androidx.datastore.preferences.preferencesDataStoreFile
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.education.dagger.ContextualEducationModule.EduDataStoreScope
+import com.android.systemui.education.data.model.GestureEduModel
+import com.android.systemui.shared.education.GestureType
+import javax.inject.Inject
+import javax.inject.Provider
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.filterNotNull
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.flow.flatMapLatest
+import kotlinx.coroutines.flow.map
+
+/**
+ * A contextual education repository to:
+ * 1) store education data per user
+ * 2) provide methods to read and update data on model-level
+ * 3) provide method to enable changing datastore when user is changed
+ */
+@SysUISingleton
+class UserContextualEducationRepository
+@Inject
+constructor(
+    @Application private val applicationContext: Context,
+    @EduDataStoreScope private val dataStoreScopeProvider: Provider<CoroutineScope>
+) {
+    companion object {
+        const val SIGNAL_COUNT_SUFFIX = "_SIGNAL_COUNT"
+        const val NUMBER_OF_EDU_SHOWN_SUFFIX = "_NUMBER_OF_EDU_SHOWN"
+
+        const val DATASTORE_DIR = "education/USER%s_ContextualEducation"
+    }
+
+    private var dataStoreScope: CoroutineScope? = null
+
+    private val datastore = MutableStateFlow<DataStore<Preferences>?>(null)
+
+    @OptIn(kotlinx.coroutines.ExperimentalCoroutinesApi::class)
+    private val prefData: Flow<Preferences> = datastore.filterNotNull().flatMapLatest { it.data }
+
+    internal fun setUser(userId: Int) {
+        dataStoreScope?.cancel()
+        val newDsScope = dataStoreScopeProvider.get()
+        datastore.value =
+            PreferenceDataStoreFactory.create(
+                produceFile = {
+                    applicationContext.preferencesDataStoreFile(
+                        String.format(DATASTORE_DIR, userId)
+                    )
+                },
+                scope = newDsScope,
+            )
+        dataStoreScope = newDsScope
+    }
+
+    internal fun readGestureEduModelFlow(gestureType: GestureType): Flow<GestureEduModel> =
+        prefData.map { preferences -> getGestureEduModel(gestureType, preferences) }
+
+    private fun getGestureEduModel(
+        gestureType: GestureType,
+        preferences: Preferences
+    ): GestureEduModel {
+        return GestureEduModel(
+            signalCount = preferences[getSignalCountKey(gestureType)] ?: 0,
+            educationShownCount = preferences[getEducationShownCountKey(gestureType)] ?: 0,
+        )
+    }
+
+    internal suspend fun updateGestureEduModel(
+        gestureType: GestureType,
+        transform: (GestureEduModel) -> GestureEduModel
+    ) {
+        datastore.filterNotNull().first().edit { preferences ->
+            val currentModel = getGestureEduModel(gestureType, preferences)
+            val updatedModel = transform(currentModel)
+            preferences[getSignalCountKey(gestureType)] = updatedModel.signalCount
+            preferences[getEducationShownCountKey(gestureType)] = updatedModel.educationShownCount
+        }
+    }
+
+    private fun getSignalCountKey(gestureType: GestureType): Preferences.Key<Int> =
+        intPreferencesKey(gestureType.name + SIGNAL_COUNT_SUFFIX)
+
+    private fun getEducationShownCountKey(gestureType: GestureType): Preferences.Key<Int> =
+        intPreferencesKey(gestureType.name + NUMBER_OF_EDU_SHOWN_SUFFIX)
+}
diff --git a/packages/SystemUI/src/com/android/systemui/inputdevice/data/repository/InputDeviceRepository.kt b/packages/SystemUI/src/com/android/systemui/inputdevice/data/repository/InputDeviceRepository.kt
new file mode 100644
index 0000000..3b161b6
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/inputdevice/data/repository/InputDeviceRepository.kt
@@ -0,0 +1,95 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.systemui.inputdevice.data.repository
+
+import android.annotation.SuppressLint
+import android.hardware.input.InputManager
+import android.os.Handler
+import com.android.systemui.common.coroutine.ChannelExt.trySendWithFailureLogging
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.utils.coroutines.flow.conflatedCallbackFlow
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.channels.SendChannel
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.shareIn
+
+@SysUISingleton
+class InputDeviceRepository
+@Inject
+constructor(
+    @Background private val backgroundHandler: Handler,
+    @Background private val backgroundScope: CoroutineScope,
+    private val inputManager: InputManager
+) {
+
+    sealed interface DeviceChange
+
+    data class DeviceAdded(val deviceId: Int) : DeviceChange
+
+    data object DeviceRemoved : DeviceChange
+
+    data object FreshStart : DeviceChange
+
+    /**
+     * Emits collection of all currently connected keyboards and what was the last [DeviceChange].
+     * It emits collection so that every new subscriber to this SharedFlow can get latest state of
+     * all keyboards. Otherwise we might get into situation where subscriber timing on
+     * initialization matter and later subscriber will only get latest device and will miss all
+     * previous devices.
+     */
+    // TODO(b/351984587): Replace with StateFlow
+    @SuppressLint("SharedFlowCreation")
+    val deviceChange: Flow<Pair<Collection<Int>, DeviceChange>> =
+        conflatedCallbackFlow {
+                var connectedDevices = inputManager.inputDeviceIds.toSet()
+                val listener =
+                    object : InputManager.InputDeviceListener {
+                        override fun onInputDeviceAdded(deviceId: Int) {
+                            connectedDevices = connectedDevices + deviceId
+                            sendWithLogging(connectedDevices to DeviceAdded(deviceId))
+                        }
+
+                        override fun onInputDeviceChanged(deviceId: Int) = Unit
+
+                        override fun onInputDeviceRemoved(deviceId: Int) {
+                            connectedDevices = connectedDevices - deviceId
+                            sendWithLogging(connectedDevices to DeviceRemoved)
+                        }
+                    }
+                sendWithLogging(connectedDevices to FreshStart)
+                inputManager.registerInputDeviceListener(listener, backgroundHandler)
+                awaitClose { inputManager.unregisterInputDeviceListener(listener) }
+            }
+            .shareIn(
+                scope = backgroundScope,
+                started = SharingStarted.Lazily,
+                replay = 1,
+            )
+
+    private fun <T> SendChannel<T>.sendWithLogging(element: T) {
+        trySendWithFailureLogging(element, TAG)
+    }
+
+    companion object {
+        const val TAG = "InputDeviceRepository"
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/data/repository/KeyboardRepository.kt b/packages/SystemUI/src/com/android/systemui/keyboard/data/repository/KeyboardRepository.kt
index 91d5280..817849c 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/data/repository/KeyboardRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/data/repository/KeyboardRepository.kt
@@ -21,21 +21,23 @@
 import android.hardware.input.InputManager.KeyboardBacklightListener
 import android.hardware.input.KeyboardBacklightState
 import com.android.systemui.common.coroutine.ChannelExt.trySendWithFailureLogging
-import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
 import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.inputdevice.data.repository.InputDeviceRepository
+import com.android.systemui.inputdevice.data.repository.InputDeviceRepository.DeviceAdded
+import com.android.systemui.inputdevice.data.repository.InputDeviceRepository.DeviceChange
+import com.android.systemui.inputdevice.data.repository.InputDeviceRepository.DeviceRemoved
+import com.android.systemui.inputdevice.data.repository.InputDeviceRepository.FreshStart
 import com.android.systemui.keyboard.data.model.Keyboard
 import com.android.systemui.keyboard.shared.model.BacklightModel
+import com.android.systemui.utils.coroutines.flow.conflatedCallbackFlow
 import java.util.concurrent.Executor
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineDispatcher
-import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.FlowPreview
 import kotlinx.coroutines.channels.SendChannel
 import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.asFlow
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.emptyFlow
@@ -44,7 +46,6 @@
 import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.mapNotNull
-import kotlinx.coroutines.flow.shareIn
 
 /**
  * Provides information about physical keyboard states. [CommandLineKeyboardRepository] can be
@@ -71,50 +72,15 @@
 class KeyboardRepositoryImpl
 @Inject
 constructor(
-    @Application private val applicationScope: CoroutineScope,
     @Background private val backgroundDispatcher: CoroutineDispatcher,
     private val inputManager: InputManager,
+    inputDeviceRepository: InputDeviceRepository
 ) : KeyboardRepository {
 
-    private sealed interface DeviceChange
-    private data class DeviceAdded(val deviceId: Int) : DeviceChange
-    private object DeviceRemoved : DeviceChange
-    private object FreshStart : DeviceChange
-
-    /**
-     * Emits collection of all currently connected keyboards and what was the last [DeviceChange].
-     * It emits collection so that every new subscriber to this SharedFlow can get latest state of
-     * all keyboards. Otherwise we might get into situation where subscriber timing on
-     * initialization matter and later subscriber will only get latest device and will miss all
-     * previous devices.
-     */
     private val keyboardsChange: Flow<Pair<Collection<Int>, DeviceChange>> =
-        conflatedCallbackFlow {
-                var connectedDevices = inputManager.inputDeviceIds.toSet()
-                val listener =
-                    object : InputManager.InputDeviceListener {
-                        override fun onInputDeviceAdded(deviceId: Int) {
-                            connectedDevices = connectedDevices + deviceId
-                            sendWithLogging(connectedDevices to DeviceAdded(deviceId))
-                        }
-
-                        override fun onInputDeviceChanged(deviceId: Int) = Unit
-
-                        override fun onInputDeviceRemoved(deviceId: Int) {
-                            connectedDevices = connectedDevices - deviceId
-                            sendWithLogging(connectedDevices to DeviceRemoved)
-                        }
-                    }
-                sendWithLogging(connectedDevices to FreshStart)
-                inputManager.registerInputDeviceListener(listener, /* handler= */ null)
-                awaitClose { inputManager.unregisterInputDeviceListener(listener) }
-            }
-            .map { (ids, change) -> ids.filter { id -> isPhysicalFullKeyboard(id) } to change }
-            .shareIn(
-                scope = applicationScope,
-                started = SharingStarted.Lazily,
-                replay = 1,
-            )
+        inputDeviceRepository.deviceChange.map { (ids, change) ->
+            ids.filter { id -> isPhysicalFullKeyboard(id) } to change
+        }
 
     @FlowPreview
     override val newlyConnectedKeyboard: Flow<Keyboard> =
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManager.java
index 0316b0e..96127b6 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManager.java
@@ -63,6 +63,7 @@
 import com.android.systemui.bouncer.ui.BouncerView;
 import com.android.systemui.dagger.SysUISingleton;
 import com.android.systemui.dagger.qualifiers.Main;
+import com.android.systemui.deviceentry.domain.interactor.DeviceEntryInteractor;
 import com.android.systemui.deviceentry.shared.DeviceEntryUdfpsRefactor;
 import com.android.systemui.dock.DockManager;
 import com.android.systemui.dreams.DreamOverlayStateController;
@@ -167,6 +168,7 @@
     private final BouncerView mPrimaryBouncerView;
     private final Lazy<ShadeController> mShadeController;
     private final Lazy<SceneInteractor> mSceneInteractorLazy;
+    private final Lazy<DeviceEntryInteractor> mDeviceEntryInteractorLazy;
 
     private Job mListenForAlternateBouncerTransitionSteps = null;
     private Job mListenForKeyguardAuthenticatedBiometricsHandled = null;
@@ -395,7 +397,8 @@
             Lazy<KeyguardSurfaceBehindInteractor> surfaceBehindInteractor,
             JavaAdapter javaAdapter,
             Lazy<SceneInteractor> sceneInteractorLazy,
-            StatusBarKeyguardViewManagerInteractor statusBarKeyguardViewManagerInteractor
+            StatusBarKeyguardViewManagerInteractor statusBarKeyguardViewManagerInteractor,
+            Lazy<DeviceEntryInteractor> deviceEntryInteractorLazy
     ) {
         mContext = context;
         mViewMediatorCallback = callback;
@@ -430,6 +433,7 @@
         mJavaAdapter = javaAdapter;
         mSceneInteractorLazy = sceneInteractorLazy;
         mStatusBarKeyguardViewManagerInteractor = statusBarKeyguardViewManagerInteractor;
+        mDeviceEntryInteractorLazy = deviceEntryInteractorLazy;
     }
 
     KeyguardTransitionInteractor mKeyguardTransitionInteractor;
@@ -735,6 +739,11 @@
      *                 {@see KeyguardBouncer#show(boolean, boolean)}
      */
     public void showBouncer(boolean scrimmed) {
+        if (SceneContainerFlag.isEnabled()) {
+            mDeviceEntryInteractorLazy.get().attemptDeviceEntry();
+            return;
+        }
+
         if (DeviceEntryUdfpsRefactor.isEnabled()) {
             if (mAlternateBouncerInteractor.canShowAlternateBouncerForFingerprint()) {
                 Log.d(TAG, "showBouncer:alternateBouncer.forceShow()");
@@ -777,8 +786,7 @@
         hideAlternateBouncer(false);
         if (mKeyguardStateController.isShowing() && !isBouncerShowing()) {
             if (SceneContainerFlag.isEnabled()) {
-                mSceneInteractorLazy.get().changeScene(
-                        Scenes.Bouncer, "StatusBarKeyguardViewManager.showPrimaryBouncer");
+                mDeviceEntryInteractorLazy.get().attemptDeviceEntry();
             } else {
                 mPrimaryBouncerInteractor.show(scrimmed);
             }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyboard/data/repository/KeyboardRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyboard/data/repository/KeyboardRepositoryTest.kt
index 53bcf86..361e768 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyboard/data/repository/KeyboardRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyboard/data/repository/KeyboardRepositoryTest.kt
@@ -20,6 +20,7 @@
 import android.hardware.input.InputManager
 import android.hardware.input.InputManager.KeyboardBacklightListener
 import android.hardware.input.KeyboardBacklightState
+import android.testing.TestableLooper
 import android.view.InputDevice
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
@@ -27,11 +28,13 @@
 import com.android.systemui.coroutines.FlowValue
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.coroutines.collectValues
+import com.android.systemui.inputdevice.data.repository.InputDeviceRepository
 import com.android.systemui.keyboard.data.model.Keyboard
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.nullable
 import com.android.systemui.util.mockito.whenever
+import com.android.systemui.utils.os.FakeHandler
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -53,6 +56,7 @@
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
+@TestableLooper.RunWithLooper
 @RunWith(AndroidJUnit4::class)
 class KeyboardRepositoryTest : SysuiTestCase() {
 
@@ -63,6 +67,7 @@
 
     private lateinit var underTest: KeyboardRepository
     private lateinit var dispatcher: CoroutineDispatcher
+    private lateinit var inputDeviceRepo: InputDeviceRepository
     private lateinit var testScope: TestScope
 
     @Before
@@ -75,7 +80,9 @@
         }
         dispatcher = StandardTestDispatcher()
         testScope = TestScope(dispatcher)
-        underTest = KeyboardRepositoryImpl(testScope.backgroundScope, dispatcher, inputManager)
+        val handler = FakeHandler(TestableLooper.get(this).looper)
+        inputDeviceRepo = InputDeviceRepository(handler, testScope.backgroundScope, inputManager)
+        underTest = KeyboardRepositoryImpl(dispatcher, inputManager, inputDeviceRepo)
     }
 
     @Test
@@ -363,6 +370,7 @@
         private val maxBrightnessLevel: Int
     ) : KeyboardBacklightState() {
         override fun getBrightnessLevel() = brightnessLevel
+
         override fun getMaxBrightnessLevel() = maxBrightnessLevel
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManagerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManagerTest.java
index 3ca4c59..0e4d892 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManagerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManagerTest.java
@@ -37,6 +37,8 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.platform.test.annotations.DisableFlags;
+import android.platform.test.annotations.EnableFlags;
 import android.platform.test.annotations.RequiresFlagsEnabled;
 import android.platform.test.flag.junit.CheckFlagsRule;
 import android.platform.test.flag.junit.DeviceFlagsValueProvider;
@@ -66,6 +68,7 @@
 import com.android.keyguard.KeyguardUpdateMonitorCallback;
 import com.android.keyguard.TrustGrantFlags;
 import com.android.keyguard.ViewMediatorCallback;
+import com.android.systemui.Flags;
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.biometrics.domain.interactor.UdfpsOverlayInteractor;
 import com.android.systemui.bouncer.domain.interactor.AlternateBouncerInteractor;
@@ -74,8 +77,11 @@
 import com.android.systemui.bouncer.domain.interactor.PrimaryBouncerInteractor;
 import com.android.systemui.bouncer.ui.BouncerView;
 import com.android.systemui.bouncer.ui.BouncerViewDelegate;
+import com.android.systemui.deviceentry.domain.interactor.DeviceEntryInteractor;
 import com.android.systemui.dock.DockManager;
 import com.android.systemui.dreams.DreamOverlayStateController;
+import com.android.systemui.flags.DisableSceneContainer;
+import com.android.systemui.flags.EnableSceneContainer;
 import com.android.systemui.keyguard.domain.interactor.KeyguardDismissActionInteractor;
 import com.android.systemui.keyguard.domain.interactor.KeyguardSurfaceBehindInteractor;
 import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor;
@@ -158,6 +164,7 @@
     @Mock private TaskbarDelegate mTaskbarDelegate;
     @Mock private StatusBarKeyguardViewManager.KeyguardViewManagerCallback mCallback;
     @Mock private SelectedUserInteractor mSelectedUserInteractor;
+    @Mock private DeviceEntryInteractor mDeviceEntryInteractor;
 
     private StatusBarKeyguardViewManager mStatusBarKeyguardViewManager;
     private PrimaryBouncerCallbackInteractor.PrimaryBouncerExpansionCallback
@@ -178,6 +185,7 @@
     public final CheckFlagsRule mCheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule();
 
     @Before
+    @DisableFlags(com.android.systemui.Flags.FLAG_KEYGUARD_WM_STATE_REFACTOR)
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         when(mContainer.findViewById(anyInt())).thenReturn(mKeyguardMessageArea);
@@ -185,10 +193,6 @@
                 .thenReturn(mKeyguardMessageAreaController);
         when(mBouncerView.getDelegate()).thenReturn(mBouncerViewDelegate);
         when(mBouncerViewDelegate.getBackCallback()).thenReturn(mBouncerViewDelegateBackCallback);
-        mSetFlagsRule.disableFlags(
-                com.android.systemui.Flags.FLAG_DEVICE_ENTRY_UDFPS_REFACTOR,
-                com.android.systemui.Flags.FLAG_KEYGUARD_WM_STATE_REFACTOR
-        );
 
         when(mNotificationShadeWindowController.getWindowRootView())
                 .thenReturn(mNotificationShadeWindowView);
@@ -227,7 +231,8 @@
                         () -> mock(KeyguardSurfaceBehindInteractor.class),
                         mock(JavaAdapter.class),
                         () -> mock(SceneInteractor.class),
-                        mock(StatusBarKeyguardViewManagerInteractor.class)) {
+                        mock(StatusBarKeyguardViewManagerInteractor.class),
+                        () -> mDeviceEntryInteractor) {
                     @Override
                     public ViewRootImpl getViewRootImpl() {
                         return mViewRootImpl;
@@ -250,6 +255,7 @@
     }
 
     @Test
+    @DisableSceneContainer
     public void dismissWithAction_AfterKeyguardGoneSetToFalse() {
         OnDismissAction action = () -> false;
         Runnable cancelAction = () -> {
@@ -265,6 +271,7 @@
         mStatusBarKeyguardViewManager.hide(0 /* startTime */, 0 /* fadeoutDuration */);
         mStatusBarKeyguardViewManager.showPrimaryBouncer(true /* scrimmed */);
         verify(mPrimaryBouncerInteractor, never()).show(anyBoolean());
+        verify(mDeviceEntryInteractor, never()).attemptDeviceEntry();
     }
 
     @Test
@@ -274,9 +281,11 @@
                 KeyguardSecurityModel.SecurityMode.Password);
         mStatusBarKeyguardViewManager.showPrimaryBouncer(true /* scrimmed */);
         verify(mPrimaryBouncerInteractor, never()).show(anyBoolean());
+        verify(mDeviceEntryInteractor, never()).attemptDeviceEntry();
     }
 
     @Test
+    @DisableSceneContainer
     public void showBouncer_showsTheBouncer() {
         mStatusBarKeyguardViewManager.showPrimaryBouncer(true /* scrimmed */);
         verify(mPrimaryBouncerInteractor).show(eq(true));
@@ -320,6 +329,7 @@
     }
 
     @Test
+    @DisableSceneContainer
     public void onPanelExpansionChanged_showsBouncerWhenSwiping() {
         mKeyguardStateController.setCanDismissLockScreen(false);
         mStatusBarKeyguardViewManager.onPanelExpansionChanged(EXPANSION_EVENT);
@@ -456,6 +466,7 @@
     }
 
     @Test
+    @DisableSceneContainer
     public void testHiding_cancelsGoneRunnable() {
         OnDismissAction action = mock(OnDismissAction.class);
         Runnable cancelAction = mock(Runnable.class);
@@ -470,6 +481,7 @@
     }
 
     @Test
+    @DisableSceneContainer
     public void testHidingBouncer_cancelsGoneRunnable() {
         OnDismissAction action = mock(OnDismissAction.class);
         Runnable cancelAction = mock(Runnable.class);
@@ -484,6 +496,7 @@
     }
 
     @Test
+    @DisableSceneContainer
     public void testHiding_doesntCancelWhenShowing() {
         OnDismissAction action = mock(OnDismissAction.class);
         Runnable cancelAction = mock(Runnable.class);
@@ -539,6 +552,7 @@
     }
 
     @Test
+    @DisableSceneContainer
     public void testShowAltAuth_unlockingWithBiometricNotAllowed() {
         // GIVEN cannot use alternate bouncer
         when(mPrimaryBouncerInteractor.isFullyShowing()).thenReturn(false);
@@ -553,6 +567,8 @@
     }
 
     @Test
+    @DisableSceneContainer
+    @DisableFlags(com.android.systemui.Flags.FLAG_DEVICE_ENTRY_UDFPS_REFACTOR)
     public void testShowAlternateBouncer_unlockingWithBiometricAllowed() {
         // GIVEN will show alternate bouncer
         when(mPrimaryBouncerInteractor.isFullyShowing()).thenReturn(false);
@@ -735,7 +751,8 @@
                         () -> mock(KeyguardSurfaceBehindInteractor.class),
                         mock(JavaAdapter.class),
                         () -> mock(SceneInteractor.class),
-                        mock(StatusBarKeyguardViewManagerInteractor.class)) {
+                        mock(StatusBarKeyguardViewManagerInteractor.class),
+                        () -> mDeviceEntryInteractor) {
                     @Override
                     public ViewRootImpl getViewRootImpl() {
                         return mViewRootImpl;
@@ -783,11 +800,11 @@
     }
 
     @Test
+    @EnableFlags(com.android.systemui.Flags.FLAG_DEVICE_ENTRY_UDFPS_REFACTOR)
     public void handleDispatchTouchEvent_alternateBouncerViewFlagEnabled() {
         mStatusBarKeyguardViewManager.addCallback(mCallback);
 
         // GIVEN alternate bouncer view flag enabled & the alternate bouncer is visible
-        mSetFlagsRule.enableFlags(com.android.systemui.Flags.FLAG_DEVICE_ENTRY_UDFPS_REFACTOR);
         when(mAlternateBouncerInteractor.isVisibleState()).thenReturn(true);
 
         // THEN the touch is not acted upon
@@ -795,9 +812,9 @@
     }
 
     @Test
+    @EnableFlags(com.android.systemui.Flags.FLAG_DEVICE_ENTRY_UDFPS_REFACTOR)
     public void onInterceptTouch_alternateBouncerViewFlagEnabled() {
         // GIVEN alternate bouncer view flag enabled & the alternate bouncer is visible
-        mSetFlagsRule.enableFlags(com.android.systemui.Flags.FLAG_DEVICE_ENTRY_UDFPS_REFACTOR);
         when(mAlternateBouncerInteractor.isVisibleState()).thenReturn(true);
 
         // THEN the touch is not intercepted
@@ -829,6 +846,8 @@
     }
 
     @Test
+    @DisableSceneContainer
+    @DisableFlags(Flags.FLAG_DEVICE_ENTRY_UDFPS_REFACTOR)
     public void handleDispatchTouchEvent_shouldInterceptTouchAndHandleTouch() {
         mStatusBarKeyguardViewManager.addCallback(mCallback);
 
@@ -855,6 +874,8 @@
     }
 
     @Test
+    @DisableSceneContainer
+    @DisableFlags(Flags.FLAG_DEVICE_ENTRY_UDFPS_REFACTOR)
     public void handleDispatchTouchEvent_shouldInterceptTouchButNotHandleTouch() {
         mStatusBarKeyguardViewManager.addCallback(mCallback);
 
@@ -881,6 +902,7 @@
     }
 
     @Test
+    @DisableSceneContainer
     public void shouldInterceptTouch_alternateBouncerNotVisible() {
         // GIVEN the alternate bouncer is not visible
         when(mAlternateBouncerInteractor.isVisibleState()).thenReturn(false);
@@ -898,6 +920,8 @@
     }
 
     @Test
+    @DisableSceneContainer
+    @DisableFlags(Flags.FLAG_DEVICE_ENTRY_UDFPS_REFACTOR)
     public void shouldInterceptTouch_alternateBouncerVisible() {
         // GIVEN the alternate bouncer is visible
         when(mAlternateBouncerInteractor.isVisibleState()).thenReturn(true);
@@ -931,6 +955,8 @@
     }
 
     @Test
+    @DisableSceneContainer
+    @DisableFlags(Flags.FLAG_DEVICE_ENTRY_UDFPS_REFACTOR)
     public void alternateBouncerOnTouch_actionDownThenUp_noMinTimeShown_noHideAltBouncer() {
         reset(mAlternateBouncerInteractor);
 
@@ -955,6 +981,8 @@
     }
 
     @Test
+    @DisableSceneContainer
+    @DisableFlags(Flags.FLAG_DEVICE_ENTRY_UDFPS_REFACTOR)
     public void alternateBouncerOnTouch_actionDownThenUp_handlesTouch_hidesAltBouncer() {
         reset(mAlternateBouncerInteractor);
 
@@ -979,6 +1007,7 @@
     }
 
     @Test
+    @DisableSceneContainer
     public void alternateBouncerOnTouch_actionUp_doesNotHideAlternateBouncer() {
         reset(mAlternateBouncerInteractor);
 
@@ -996,6 +1025,8 @@
     }
 
     @Test
+    @DisableSceneContainer
+    @DisableFlags(Flags.FLAG_DEVICE_ENTRY_UDFPS_REFACTOR)
     public void onTrustChanged_hideAlternateBouncerAndClearMessageArea() {
         // GIVEN keyguard update monitor callback is registered
         verify(mKeyguardUpdateMonitor).registerCallback(mKeyguardUpdateMonitorCallback.capture());
@@ -1024,6 +1055,7 @@
     }
 
     @Test
+    @DisableSceneContainer
     public void testShowBouncerOrKeyguard_needsFullScreen() {
         when(mKeyguardSecurityModel.getSecurityMode(anyInt())).thenReturn(
                 KeyguardSecurityModel.SecurityMode.SimPin);
@@ -1033,6 +1065,7 @@
     }
 
     @Test
+    @DisableSceneContainer
     public void testShowBouncerOrKeyguard_needsFullScreen_bouncerAlreadyShowing() {
         when(mKeyguardSecurityModel.getSecurityMode(anyInt())).thenReturn(
                 KeyguardSecurityModel.SecurityMode.SimPin);
@@ -1043,6 +1076,20 @@
     }
 
     @Test
+    @EnableSceneContainer
+    public void showBouncer_attemptDeviceEntry() {
+        mStatusBarKeyguardViewManager.showBouncer(false);
+        verify(mDeviceEntryInteractor).attemptDeviceEntry();
+    }
+
+    @Test
+    @EnableSceneContainer
+    public void showPrimaryBouncer_attemptDeviceEntry() {
+        mStatusBarKeyguardViewManager.showPrimaryBouncer(false);
+        verify(mDeviceEntryInteractor).attemptDeviceEntry();
+    }
+
+    @Test
     public void altBouncerNotVisible_keyguardAuthenticatedBiometricsHandled() {
         clearInvocations(mAlternateBouncerInteractor);
         when(mAlternateBouncerInteractor.isVisibleState()).thenReturn(false);
diff --git a/services/tests/wmtests/src/com/android/server/wm/ProtoLogIntegrationTest.java b/services/tests/wmtests/src/com/android/server/wm/ProtoLogIntegrationTest.java
index 7efbc88..35cff7a 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ProtoLogIntegrationTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ProtoLogIntegrationTest.java
@@ -25,11 +25,11 @@
 
 import androidx.test.filters.SmallTest;
 
+import com.android.internal.protolog.ProtoLog;
 import com.android.internal.protolog.ProtoLogGroup;
 import com.android.internal.protolog.ProtoLogImpl;
 import com.android.internal.protolog.common.IProtoLog;
 import com.android.internal.protolog.common.LogLevel;
-import com.android.internal.protolog.ProtoLog;
 
 import org.junit.After;
 import org.junit.Ignore;
@@ -53,9 +53,6 @@
         runWith(mockedProtoLog, this::testProtoLog);
         verify(mockedProtoLog).log(eq(LogLevel.ERROR), eq(ProtoLogGroup.TEST_GROUP),
                 anyInt(), eq(0b0010010111),
-                eq(com.android.internal.protolog.ProtoLogGroup.TEST_GROUP.isLogToLogcat()
-                        ? "Test completed successfully: %b %d %x %f %% %s"
-                        : null),
                 eq(new Object[]{true, 1L, 2L, 0.3, "ok"}));
     }
 
diff --git a/tests/Internal/src/com/android/internal/protolog/LegacyProtoLogImplTest.java b/tests/Internal/src/com/android/internal/protolog/LegacyProtoLogImplTest.java
index 5a27593..5a48327 100644
--- a/tests/Internal/src/com/android/internal/protolog/LegacyProtoLogImplTest.java
+++ b/tests/Internal/src/com/android/internal/protolog/LegacyProtoLogImplTest.java
@@ -59,7 +59,6 @@
 import java.io.InputStream;
 import java.io.PrintWriter;
 import java.util.LinkedList;
-import java.util.TreeMap;
 
 /**
  * Test class for {@link ProtoLogImpl}.
@@ -90,7 +89,7 @@
         //noinspection ResultOfMethodCallIgnored
         mFile.delete();
         mProtoLog = new LegacyProtoLogImpl(mFile, mViewerConfigFilename,
-                1024 * 1024, mReader, 1024, new TreeMap<>(), () -> {});
+                1024 * 1024, mReader, 1024, () -> {});
     }
 
     @After
@@ -142,7 +141,7 @@
         TestProtoLogGroup.TEST_GROUP.setLogToProto(false);
 
         implSpy.log(
-                LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 1234, 4321, null,
+                LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 1234, 4321,
                 new Object[]{true, 10000, 30000, "test", 0.000003});
 
         verify(implSpy).passToLogcat(eq(TestProtoLogGroup.TEST_GROUP.getTag()), eq(
@@ -159,7 +158,7 @@
         TestProtoLogGroup.TEST_GROUP.setLogToProto(false);
 
         implSpy.log(
-                LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 1234, 4321, null,
+                LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 1234, 4321,
                 new Object[]{true, 10000, 0.0001, 0.00002, "test"});
 
         verify(implSpy).passToLogcat(eq(TestProtoLogGroup.TEST_GROUP.getTag()), eq(
@@ -176,7 +175,7 @@
         TestProtoLogGroup.TEST_GROUP.setLogToProto(false);
 
         implSpy.log(
-                LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 1234, 4321, "test %d",
+                LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 1234, 4321,
                 new Object[]{5});
 
         verify(implSpy).passToLogcat(eq(TestProtoLogGroup.TEST_GROUP.getTag()), eq(
@@ -192,7 +191,7 @@
         TestProtoLogGroup.TEST_GROUP.setLogToProto(false);
 
         implSpy.log(
-                LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 1234, 4321, null,
+                LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 1234, 4321,
                 new Object[]{5});
 
         verify(implSpy).passToLogcat(eq(TestProtoLogGroup.TEST_GROUP.getTag()), eq(
@@ -208,7 +207,7 @@
         TestProtoLogGroup.TEST_GROUP.setLogToProto(false);
 
         implSpy.log(
-                LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 1234, 4321, "test %d",
+                LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 1234, 4321,
                 new Object[]{5});
 
         verify(implSpy, never()).passToLogcat(any(), any(), any());
@@ -277,7 +276,7 @@
         long before = SystemClock.elapsedRealtimeNanos();
         mProtoLog.log(
                 LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 1234,
-                0b1110101001010100, null,
+                0b1110101001010100,
                 new Object[]{"test", 1, 2, 3, 0.4, 0.5, 0.6, true});
         long after = SystemClock.elapsedRealtimeNanos();
         mProtoLog.stopProtoLog(mock(PrintWriter.class), true);
@@ -302,7 +301,7 @@
         long before = SystemClock.elapsedRealtimeNanos();
         mProtoLog.log(
                 LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 1234,
-                0b01100100, null,
+                0b01100100,
                 new Object[]{"test", 1, 0.1, true});
         long after = SystemClock.elapsedRealtimeNanos();
         mProtoLog.stopProtoLog(mock(PrintWriter.class), true);
@@ -326,7 +325,7 @@
         TestProtoLogGroup.TEST_GROUP.setLogToProto(false);
         mProtoLog.startProtoLog(mock(PrintWriter.class));
         mProtoLog.log(LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 1234,
-                0b11, null, new Object[]{true});
+                0b11, new Object[]{true});
         mProtoLog.stopProtoLog(mock(PrintWriter.class), true);
         try (InputStream is = new FileInputStream(mFile)) {
             ProtoInputStream ip = new ProtoInputStream(is);
diff --git a/tests/Internal/src/com/android/internal/protolog/PerfettoProtoLogImplTest.java b/tests/Internal/src/com/android/internal/protolog/PerfettoProtoLogImplTest.java
index 1d7b6b3..b6672a0 100644
--- a/tests/Internal/src/com/android/internal/protolog/PerfettoProtoLogImplTest.java
+++ b/tests/Internal/src/com/android/internal/protolog/PerfettoProtoLogImplTest.java
@@ -60,6 +60,9 @@
 import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 
+import perfetto.protos.Protolog;
+import perfetto.protos.ProtologCommon;
+
 import java.io.File;
 import java.io.IOException;
 import java.util.List;
@@ -67,9 +70,6 @@
 import java.util.TreeMap;
 import java.util.concurrent.atomic.AtomicInteger;
 
-import perfetto.protos.Protolog;
-import perfetto.protos.ProtologCommon;
-
 /**
  * Test class for {@link ProtoLogImpl}.
  */
@@ -111,6 +111,9 @@
         //noinspection ResultOfMethodCallIgnored
         mFile.delete();
 
+        TestProtoLogGroup.TEST_GROUP.setLogToLogcat(false);
+        TestProtoLogGroup.TEST_GROUP.setLogToProto(false);
+
         mViewerConfigBuilder = Protolog.ProtoLogViewerConfig.newBuilder()
                 .addGroups(
                         Protolog.ProtoLogViewerConfig.Group.newBuilder()
@@ -157,8 +160,9 @@
         mCacheUpdater = () -> {};
         mReader = Mockito.spy(new ProtoLogViewerConfigReader(viewerConfigInputStreamProvider));
         mProtoLog = new PerfettoProtoLogImpl(
-                viewerConfigInputStreamProvider, mReader, new TreeMap<>(),
+                viewerConfigInputStreamProvider, mReader,
                 () -> mCacheUpdater.run());
+        mProtoLog.registerGroups(TestProtoLogGroup.values());
     }
 
     @After
@@ -210,15 +214,15 @@
             // Shouldn't be logging anything except WTF unless explicitly requested in the group
             // override.
             mProtoLog.log(LogLevel.DEBUG, TestProtoLogGroup.TEST_GROUP, 1,
-                    LogDataType.BOOLEAN, null, new Object[]{true});
+                    LogDataType.BOOLEAN, new Object[]{true});
             mProtoLog.log(LogLevel.VERBOSE, TestProtoLogGroup.TEST_GROUP, 2,
-                    LogDataType.BOOLEAN, null, new Object[]{true});
+                    LogDataType.BOOLEAN, new Object[]{true});
             mProtoLog.log(LogLevel.WARN, TestProtoLogGroup.TEST_GROUP, 3,
-                    LogDataType.BOOLEAN, null, new Object[]{true});
+                    LogDataType.BOOLEAN, new Object[]{true});
             mProtoLog.log(LogLevel.ERROR, TestProtoLogGroup.TEST_GROUP, 4,
-                    LogDataType.BOOLEAN, null, new Object[]{true});
+                    LogDataType.BOOLEAN, new Object[]{true});
             mProtoLog.log(LogLevel.WTF, TestProtoLogGroup.TEST_GROUP, 5,
-                    LogDataType.BOOLEAN, null, new Object[]{true});
+                    LogDataType.BOOLEAN, new Object[]{true});
         } finally {
             traceMonitor.stop(mWriter);
         }
@@ -240,15 +244,15 @@
         try {
             traceMonitor.start();
             mProtoLog.log(LogLevel.DEBUG, TestProtoLogGroup.TEST_GROUP, 1,
-                    LogDataType.BOOLEAN, null, new Object[]{true});
+                    LogDataType.BOOLEAN, new Object[]{true});
             mProtoLog.log(LogLevel.VERBOSE, TestProtoLogGroup.TEST_GROUP, 2,
-                    LogDataType.BOOLEAN, null, new Object[]{true});
+                    LogDataType.BOOLEAN, new Object[]{true});
             mProtoLog.log(LogLevel.WARN, TestProtoLogGroup.TEST_GROUP, 3,
-                    LogDataType.BOOLEAN, null, new Object[]{true});
+                    LogDataType.BOOLEAN, new Object[]{true});
             mProtoLog.log(LogLevel.ERROR, TestProtoLogGroup.TEST_GROUP, 4,
-                    LogDataType.BOOLEAN, null, new Object[]{true});
+                    LogDataType.BOOLEAN, new Object[]{true});
             mProtoLog.log(LogLevel.WTF, TestProtoLogGroup.TEST_GROUP, 5,
-                    LogDataType.BOOLEAN, null, new Object[]{true});
+                    LogDataType.BOOLEAN, new Object[]{true});
         } finally {
             traceMonitor.stop(mWriter);
         }
@@ -274,15 +278,15 @@
         try {
             traceMonitor.start();
             mProtoLog.log(LogLevel.DEBUG, TestProtoLogGroup.TEST_GROUP, 1,
-                    LogDataType.BOOLEAN, null, new Object[]{true});
+                    LogDataType.BOOLEAN, new Object[]{true});
             mProtoLog.log(LogLevel.VERBOSE, TestProtoLogGroup.TEST_GROUP, 2,
-                    LogDataType.BOOLEAN, null, new Object[]{true});
+                    LogDataType.BOOLEAN, new Object[]{true});
             mProtoLog.log(LogLevel.WARN, TestProtoLogGroup.TEST_GROUP, 3,
-                    LogDataType.BOOLEAN, null, new Object[]{true});
+                    LogDataType.BOOLEAN, new Object[]{true});
             mProtoLog.log(LogLevel.ERROR, TestProtoLogGroup.TEST_GROUP, 4,
-                    LogDataType.BOOLEAN, null, new Object[]{true});
+                    LogDataType.BOOLEAN, new Object[]{true});
             mProtoLog.log(LogLevel.WTF, TestProtoLogGroup.TEST_GROUP, 5,
-                    LogDataType.BOOLEAN, null, new Object[]{true});
+                    LogDataType.BOOLEAN, new Object[]{true});
         } finally {
             traceMonitor.stop(mWriter);
         }
@@ -304,15 +308,15 @@
         try {
             traceMonitor.start();
             mProtoLog.log(LogLevel.DEBUG, TestProtoLogGroup.TEST_GROUP, 1,
-                    LogDataType.BOOLEAN, null, new Object[]{true});
+                    LogDataType.BOOLEAN, new Object[]{true});
             mProtoLog.log(LogLevel.VERBOSE, TestProtoLogGroup.TEST_GROUP, 2,
-                    LogDataType.BOOLEAN, null, new Object[]{true});
+                    LogDataType.BOOLEAN, new Object[]{true});
             mProtoLog.log(LogLevel.WARN, TestProtoLogGroup.TEST_GROUP, 3,
-                    LogDataType.BOOLEAN, null, new Object[]{true});
+                    LogDataType.BOOLEAN, new Object[]{true});
             mProtoLog.log(LogLevel.ERROR, TestProtoLogGroup.TEST_GROUP, 4,
-                    LogDataType.BOOLEAN, null, new Object[]{true});
+                    LogDataType.BOOLEAN, new Object[]{true});
             mProtoLog.log(LogLevel.WTF, TestProtoLogGroup.TEST_GROUP, 5,
-                    LogDataType.BOOLEAN, null, new Object[]{true});
+                    LogDataType.BOOLEAN, new Object[]{true});
         } finally {
             traceMonitor.stop(mWriter);
         }
@@ -329,14 +333,14 @@
     }
 
     @Test
-    public void log_logcatEnabledExternalMessage() {
+    public void log_logcatEnabled() {
         when(mReader.getViewerString(anyLong())).thenReturn("test %b %d %% 0x%x %s %f");
         PerfettoProtoLogImpl implSpy = Mockito.spy(mProtoLog);
         TestProtoLogGroup.TEST_GROUP.setLogToLogcat(true);
         TestProtoLogGroup.TEST_GROUP.setLogToProto(false);
 
         implSpy.log(
-                LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 1234, 4321, null,
+                LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 1234, 4321,
                 new Object[]{true, 10000, 30000, "test", 0.000003});
 
         verify(implSpy).passToLogcat(eq(TestProtoLogGroup.TEST_GROUP.getTag()), eq(
@@ -353,32 +357,17 @@
         TestProtoLogGroup.TEST_GROUP.setLogToProto(false);
 
         implSpy.log(
-                LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 1234, 4321, null,
+                LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 1234, 4321,
                 new Object[]{true, 10000, 0.0001, 0.00002, "test"});
 
         verify(implSpy).passToLogcat(eq(TestProtoLogGroup.TEST_GROUP.getTag()), eq(
                 LogLevel.INFO),
-                eq("UNKNOWN MESSAGE (1234) true 10000 1.0E-4 2.0E-5 test"));
+                eq("FORMAT_ERROR \"test %b %d %% %x %s %f\", "
+                        + "args=(true, 10000, 1.0E-4, 2.0E-5, test)"));
         verify(mReader).getViewerString(eq(1234L));
     }
 
     @Test
-    public void log_logcatEnabledInlineMessage() {
-        when(mReader.getViewerString(anyLong())).thenReturn("test %d");
-        PerfettoProtoLogImpl implSpy = Mockito.spy(mProtoLog);
-        TestProtoLogGroup.TEST_GROUP.setLogToLogcat(true);
-        TestProtoLogGroup.TEST_GROUP.setLogToProto(false);
-
-        implSpy.log(
-                LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 1234, 4321, "test %d",
-                new Object[]{5});
-
-        verify(implSpy).passToLogcat(eq(TestProtoLogGroup.TEST_GROUP.getTag()), eq(
-                LogLevel.INFO), eq("test 5"));
-        verify(mReader, never()).getViewerString(anyLong());
-    }
-
-    @Test
     public void log_logcatEnabledNoMessage() {
         when(mReader.getViewerString(anyLong())).thenReturn(null);
         PerfettoProtoLogImpl implSpy = Mockito.spy(mProtoLog);
@@ -386,11 +375,11 @@
         TestProtoLogGroup.TEST_GROUP.setLogToProto(false);
 
         implSpy.log(
-                LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 1234, 4321, null,
+                LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 1234, 4321,
                 new Object[]{5});
 
         verify(implSpy).passToLogcat(eq(TestProtoLogGroup.TEST_GROUP.getTag()), eq(
-                LogLevel.INFO), eq("UNKNOWN MESSAGE (1234) 5"));
+                LogLevel.INFO), eq("UNKNOWN MESSAGE#1234 (5)"));
         verify(mReader).getViewerString(eq(1234L));
     }
 
@@ -401,7 +390,7 @@
         TestProtoLogGroup.TEST_GROUP.setLogToLogcat(false);
 
         implSpy.log(
-                LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 1234, 4321, "test %d",
+                LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 1234, 4321,
                 new Object[]{5});
 
         verify(implSpy, never()).passToLogcat(any(), any(), any());
@@ -425,7 +414,7 @@
             before = SystemClock.elapsedRealtimeNanos();
             mProtoLog.log(
                     LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, messageHash,
-                    0b1110101001010100, null,
+                    0b1110101001010100,
                     new Object[]{"test", 1, 2, 3, 0.4, 0.5, 0.6, true});
             after = SystemClock.elapsedRealtimeNanos();
         } finally {
@@ -444,6 +433,38 @@
                 .isEqualTo("My test message :: test, 2, 4, 6, 0.400000, 5.000000e-01, 0.6, true");
     }
 
+    @Test
+    public void log_noProcessing() throws IOException {
+        PerfettoTraceMonitor traceMonitor =
+                PerfettoTraceMonitor.newBuilder().enableProtoLog().build();
+        long before;
+        long after;
+        try {
+            traceMonitor.start();
+            assertTrue(mProtoLog.isProtoEnabled());
+
+            before = SystemClock.elapsedRealtimeNanos();
+            mProtoLog.log(
+                    LogLevel.INFO, TestProtoLogGroup.TEST_GROUP,
+                    "My test message :: %s, %d, %o, %x, %f, %b",
+                    "test", 1, 2, 3, 0.4, true);
+            after = SystemClock.elapsedRealtimeNanos();
+        } finally {
+            traceMonitor.stop(mWriter);
+        }
+
+        final ResultReader reader = new ResultReader(mWriter.write(), mTraceConfig);
+        final ProtoLogTrace protolog = reader.readProtoLogTrace();
+
+        Truth.assertThat(protolog.messages).hasSize(1);
+        Truth.assertThat(protolog.messages.getFirst().getTimestamp().getElapsedNanos())
+                .isAtLeast(before);
+        Truth.assertThat(protolog.messages.getFirst().getTimestamp().getElapsedNanos())
+                .isAtMost(after);
+        Truth.assertThat(protolog.messages.getFirst().getMessage())
+                .isEqualTo("My test message :: test, 2, 4, 6, 0.400000, true");
+    }
+
     private long addMessageToConfig(ProtologCommon.ProtoLogLevel logLevel, String message) {
         final long messageId = new Random().nextLong();
         mViewerConfigBuilder.addMessages(Protolog.ProtoLogViewerConfig.MessageData.newBuilder()
@@ -470,7 +491,7 @@
             before = SystemClock.elapsedRealtimeNanos();
             mProtoLog.log(
                     LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, messageHash,
-                    0b01100100, null,
+                    0b01100100,
                     new Object[]{"test", 1, 0.1, true});
             after = SystemClock.elapsedRealtimeNanos();
         } finally {
@@ -488,7 +509,7 @@
         try {
             traceMonitor.start();
             mProtoLog.log(LogLevel.DEBUG, TestProtoLogGroup.TEST_GROUP, 1,
-                    0b11, null, new Object[]{true});
+                    0b11, new Object[]{true});
         } finally {
             traceMonitor.stop(mWriter);
         }
@@ -512,7 +533,7 @@
 
             ProtoLogImpl.setSingleInstance(mProtoLog);
             ProtoLogImpl.d(TestProtoLogGroup.TEST_GROUP, 1,
-                    0b11, null, true);
+                    0b11, true);
         } finally {
             traceMonitor.stop(mWriter);
         }
@@ -586,7 +607,7 @@
                 .isFalse();
         Truth.assertThat(mProtoLog.isEnabled(TestProtoLogGroup.TEST_GROUP, LogLevel.ERROR))
                 .isFalse();
-        Truth.assertThat(mProtoLog.isEnabled(TestProtoLogGroup.TEST_GROUP, LogLevel.WTF)).isTrue();
+        Truth.assertThat(mProtoLog.isEnabled(TestProtoLogGroup.TEST_GROUP, LogLevel.WTF)).isFalse();
 
         PerfettoTraceMonitor traceMonitor1 =
                 PerfettoTraceMonitor.newBuilder().enableProtoLog(true,
@@ -664,7 +685,53 @@
         Truth.assertThat(mProtoLog.isEnabled(TestProtoLogGroup.TEST_GROUP, LogLevel.ERROR))
                 .isFalse();
         Truth.assertThat(mProtoLog.isEnabled(TestProtoLogGroup.TEST_GROUP, LogLevel.WTF))
-                .isTrue();
+                .isFalse();
+    }
+
+    @Test
+    public void supportsNullString() throws IOException {
+        PerfettoTraceMonitor traceMonitor =
+                PerfettoTraceMonitor.newBuilder().enableProtoLog(true)
+                        .build();
+
+        try {
+            traceMonitor.start();
+
+            mProtoLog.log(LogLevel.DEBUG, TestProtoLogGroup.TEST_GROUP,
+                    "My test null string: %s", null);
+        } finally {
+            traceMonitor.stop(mWriter);
+        }
+
+        final ResultReader reader = new ResultReader(mWriter.write(), mTraceConfig);
+        final ProtoLogTrace protolog = reader.readProtoLogTrace();
+
+        Truth.assertThat(protolog.messages).hasSize(1);
+        Truth.assertThat(protolog.messages.get(0).getMessage())
+                .isEqualTo("My test null string: null");
+    }
+
+    @Test
+    public void supportNullParams() throws IOException {
+        PerfettoTraceMonitor traceMonitor =
+                PerfettoTraceMonitor.newBuilder().enableProtoLog(true)
+                        .build();
+
+        try {
+            traceMonitor.start();
+
+            mProtoLog.log(LogLevel.DEBUG, TestProtoLogGroup.TEST_GROUP,
+                    "My null args: %d, %f, %b", null, null, null);
+        } finally {
+            traceMonitor.stop(mWriter);
+        }
+
+        final ResultReader reader = new ResultReader(mWriter.write(), mTraceConfig);
+        final ProtoLogTrace protolog = reader.readProtoLogTrace();
+
+        Truth.assertThat(protolog.messages).hasSize(1);
+        Truth.assertThat(protolog.messages.get(0).getMessage())
+                .isEqualTo("My null args: 0, 0, false");
     }
 
     private enum TestProtoLogGroup implements IProtoLogGroup {
diff --git a/tests/Internal/src/com/android/internal/protolog/ProtoLogImplTest.java b/tests/Internal/src/com/android/internal/protolog/ProtoLogImplTest.java
index 60456f9..0496240 100644
--- a/tests/Internal/src/com/android/internal/protolog/ProtoLogImplTest.java
+++ b/tests/Internal/src/com/android/internal/protolog/ProtoLogImplTest.java
@@ -58,51 +58,50 @@
     public void d_logCalled() {
         IProtoLog mockedProtoLog = mock(IProtoLog.class);
         ProtoLogImpl.setSingleInstance(mockedProtoLog);
-        ProtoLogImpl.d(TestProtoLogGroup.TEST_GROUP, 1234, 4321, "test %d");
+        ProtoLogImpl.d(TestProtoLogGroup.TEST_GROUP, 1234, 4321);
         verify(mockedProtoLog).log(eq(LogLevel.DEBUG), eq(
                 TestProtoLogGroup.TEST_GROUP),
-                eq(1234L), eq(4321), eq("test %d"), eq(new Object[]{}));
+                eq(1234L), eq(4321), eq(new Object[]{}));
     }
 
     @Test
     public void v_logCalled() {
         IProtoLog mockedProtoLog = mock(IProtoLog.class);
         ProtoLogImpl.setSingleInstance(mockedProtoLog);
-        ProtoLogImpl.v(TestProtoLogGroup.TEST_GROUP, 1234, 4321, "test %d");
+        ProtoLogImpl.v(TestProtoLogGroup.TEST_GROUP, 1234, 4321);
         verify(mockedProtoLog).log(eq(LogLevel.VERBOSE), eq(
                 TestProtoLogGroup.TEST_GROUP),
-                eq(1234L), eq(4321), eq("test %d"), eq(new Object[]{}));
+                eq(1234L), eq(4321), eq(new Object[]{}));
     }
 
     @Test
     public void i_logCalled() {
         IProtoLog mockedProtoLog = mock(IProtoLog.class);
         ProtoLogImpl.setSingleInstance(mockedProtoLog);
-        ProtoLogImpl.i(TestProtoLogGroup.TEST_GROUP, 1234, 4321, "test %d");
+        ProtoLogImpl.i(TestProtoLogGroup.TEST_GROUP, 1234, 4321);
         verify(mockedProtoLog).log(eq(LogLevel.INFO), eq(
                 TestProtoLogGroup.TEST_GROUP),
-                eq(1234L), eq(4321), eq("test %d"), eq(new Object[]{}));
+                eq(1234L), eq(4321), eq(new Object[]{}));
     }
 
     @Test
     public void w_logCalled() {
         IProtoLog mockedProtoLog = mock(IProtoLog.class);
         ProtoLogImpl.setSingleInstance(mockedProtoLog);
-        ProtoLogImpl.w(TestProtoLogGroup.TEST_GROUP, 1234,
-                4321, "test %d");
+        ProtoLogImpl.w(TestProtoLogGroup.TEST_GROUP, 1234, 4321);
         verify(mockedProtoLog).log(eq(LogLevel.WARN), eq(
                 TestProtoLogGroup.TEST_GROUP),
-                eq(1234L), eq(4321), eq("test %d"), eq(new Object[]{}));
+                eq(1234L), eq(4321), eq(new Object[]{}));
     }
 
     @Test
     public void e_logCalled() {
         IProtoLog mockedProtoLog = mock(IProtoLog.class);
         ProtoLogImpl.setSingleInstance(mockedProtoLog);
-        ProtoLogImpl.e(TestProtoLogGroup.TEST_GROUP, 1234, 4321, "test %d");
+        ProtoLogImpl.e(TestProtoLogGroup.TEST_GROUP, 1234, 4321);
         verify(mockedProtoLog).log(eq(LogLevel.ERROR), eq(
                 TestProtoLogGroup.TEST_GROUP),
-                eq(1234L), eq(4321), eq("test %d"), eq(new Object[]{}));
+                eq(1234L), eq(4321), eq(new Object[]{}));
     }
 
     @Test
@@ -110,10 +109,10 @@
         IProtoLog mockedProtoLog = mock(IProtoLog.class);
         ProtoLogImpl.setSingleInstance(mockedProtoLog);
         ProtoLogImpl.wtf(TestProtoLogGroup.TEST_GROUP,
-                1234, 4321, "test %d");
+                1234, 4321);
         verify(mockedProtoLog).log(eq(LogLevel.WTF), eq(
                 TestProtoLogGroup.TEST_GROUP),
-                eq(1234L), eq(4321), eq("test %d"), eq(new Object[]{}));
+                eq(1234L), eq(4321), eq(new Object[]{}));
     }
 
     private enum TestProtoLogGroup implements IProtoLogGroup {
diff --git a/tools/protologtool/src/com/android/protolog/tool/ProtoLogCallProcessorImpl.kt b/tools/protologtool/src/com/android/protolog/tool/ProtoLogCallProcessorImpl.kt
index 1087ae6..3c99e68 100644
--- a/tools/protologtool/src/com/android/protolog/tool/ProtoLogCallProcessorImpl.kt
+++ b/tools/protologtool/src/com/android/protolog/tool/ProtoLogCallProcessorImpl.kt
@@ -120,6 +120,8 @@
 
                         logCallVisitor?.processCall(call, messageString, getLevelForMethodName(
                             call.name.toString(), call, context), groupMap.getValue(groupName))
+                    } else if (call.name.id == "initialize") {
+                        // No processing
                     } else {
                         // Process non-log message calls
                         otherCallVisitor?.processCall(call)
diff --git a/tools/protologtool/src/com/android/protolog/tool/ProtoLogTool.kt b/tools/protologtool/src/com/android/protolog/tool/ProtoLogTool.kt
index 2d5b50b..aa53005 100644
--- a/tools/protologtool/src/com/android/protolog/tool/ProtoLogTool.kt
+++ b/tools/protologtool/src/com/android/protolog/tool/ProtoLogTool.kt
@@ -443,10 +443,10 @@
             val command = CommandOptions(args)
             invoke(command)
         } catch (ex: InvalidCommandException) {
-            println("\n${ex.message}\n")
+            println("InvalidCommandException: \n${ex.message}\n")
             showHelpAndExit()
         } catch (ex: CodeProcessingException) {
-            println("\n${ex.message}\n")
+            println("CodeProcessingException: \n${ex.message}\n")
             exitProcess(1)
         }
     }
diff --git a/tools/protologtool/src/com/android/protolog/tool/SourceTransformer.kt b/tools/protologtool/src/com/android/protolog/tool/SourceTransformer.kt
index 6a8a071..c478f58 100644
--- a/tools/protologtool/src/com/android/protolog/tool/SourceTransformer.kt
+++ b/tools/protologtool/src/com/android/protolog/tool/SourceTransformer.kt
@@ -130,28 +130,27 @@
         val hash = CodeUtils.hash(packagePath, messageString, level, group)
 
         val newCall = call.clone()
-        if (!group.textEnabled) {
-            // Remove message string if text logging is not enabled by default.
-            // Out: ProtoLog.e(GROUP, null, arg)
-            newCall.arguments[1].replace(NameExpr("null"))
-        }
+        // Remove message string.
+        // Out: ProtoLog.e(GROUP, args)
+        newCall.arguments.removeAt(1)
         // Insert message string hash as a second argument.
-        // Out: ProtoLog.e(GROUP, 1234, null, arg)
+        // Out: ProtoLog.e(GROUP, 1234, args)
         newCall.arguments.add(1, LongLiteralExpr("" + hash + "L"))
         val argTypes = LogDataType.parseFormatString(messageString)
         val typeMask = LogDataType.logDataTypesToBitMask(argTypes)
         // Insert bitmap representing which Number parameters are to be considered as
         // floating point numbers.
-        // Out: ProtoLog.e(GROUP, 1234, 0, null, arg)
+        // Out: ProtoLog.e(GROUP, 1234, 0, args)
         newCall.arguments.add(2, IntegerLiteralExpr(typeMask))
         // Replace call to a stub method with an actual implementation.
-        // Out: ProtoLogImpl.e(GROUP, 1234, null, arg)
+        // Out: ProtoLogImpl.e(GROUP, 1234, 0, args)
         newCall.setScope(protoLogImplClassNode)
         if (argTypes.size != call.arguments.size - 2) {
             throw InvalidProtoLogCallException(
                 "Number of arguments (${argTypes.size} does not match format" +
                         " string in: $call", ParsingContext(path, call))
         }
+        val argsOffset = 3
         val blockStmt = BlockStmt()
         if (argTypes.isNotEmpty()) {
             // Assign every argument to a variable to check its type in compile time
@@ -160,9 +159,9 @@
             argTypes.forEachIndexed { idx, type ->
                 val varName = "protoLogParam$idx"
                 val declaration = VariableDeclarator(getASTTypeForDataType(type), varName,
-                    getConversionForType(type)(newCall.arguments[idx + 4].clone()))
+                    getConversionForType(type)(newCall.arguments[idx + argsOffset].clone()))
                 blockStmt.addStatement(ExpressionStmt(VariableDeclarationExpr(declaration)))
-                newCall.setArgument(idx + 4, NameExpr(SimpleName(varName)))
+                newCall.setArgument(idx + argsOffset, NameExpr(SimpleName(varName)))
             }
         } else {
             // Assign (Object[])null as the vararg parameter to prevent allocating an empty
diff --git a/tools/protologtool/tests/com/android/protolog/tool/EndToEndTest.kt b/tools/protologtool/tests/com/android/protolog/tool/EndToEndTest.kt
index 2a83677..0cbbd48 100644
--- a/tools/protologtool/tests/com/android/protolog/tool/EndToEndTest.kt
+++ b/tools/protologtool/tests/com/android/protolog/tool/EndToEndTest.kt
@@ -60,7 +60,7 @@
                 .containsMatch(Pattern.compile("\\{ String protoLogParam0 = " +
                         "String\\.valueOf\\(argString\\); long protoLogParam1 = argInt; " +
                         "com\\.android\\.internal\\.protolog.ProtoLogImpl_.*\\.d\\(" +
-                        "GROUP, -6872339441335321086L, 4, null, protoLogParam0, protoLogParam1" +
+                        "GROUP, -6872339441335321086L, 4, protoLogParam0, protoLogParam1" +
                         "\\); \\}"))
     }
 
diff --git a/tools/protologtool/tests/com/android/protolog/tool/SourceTransformerTest.kt b/tools/protologtool/tests/com/android/protolog/tool/SourceTransformerTest.kt
index 82aa93d..6cde7a7 100644
--- a/tools/protologtool/tests/com/android/protolog/tool/SourceTransformerTest.kt
+++ b/tools/protologtool/tests/com/android/protolog/tool/SourceTransformerTest.kt
@@ -76,7 +76,7 @@
 
             class Test {
                 void test() {
-                    if (org.example.ProtoLogImpl.Cache.TEST_GROUP_enabled[3]) { long protoLogParam0 = 100; double protoLogParam1 = 0.1; org.example.ProtoLogImpl.w(TEST_GROUP, -1473209266730422156L, 9, "test %d %f", protoLogParam0, protoLogParam1); }
+                    if (org.example.ProtoLogImpl.Cache.TEST_GROUP_enabled[3]) { long protoLogParam0 = 100; double protoLogParam1 = 0.1; org.example.ProtoLogImpl.w(TEST_GROUP, -1473209266730422156L, 9, protoLogParam0, protoLogParam1); }
                 }
             }
             """.trimIndent()
@@ -86,7 +86,7 @@
 
             class Test {
                 void test() {
-                    if (org.example.ProtoLogImpl.Cache.TEST_GROUP_enabled[3]) { long protoLogParam0 = 100; double protoLogParam1 = 0.1; String protoLogParam2 = String.valueOf("test"); org.example.ProtoLogImpl.w(TEST_GROUP, -4447034859795564700L, 9, "test %d %f " + "abc %s\n test", protoLogParam0, protoLogParam1, protoLogParam2); 
+                    if (org.example.ProtoLogImpl.Cache.TEST_GROUP_enabled[3]) { long protoLogParam0 = 100; double protoLogParam1 = 0.1; String protoLogParam2 = String.valueOf("test"); org.example.ProtoLogImpl.w(TEST_GROUP, -4447034859795564700L, 9, protoLogParam0, protoLogParam1, protoLogParam2); 
             
             }
                 }
@@ -98,8 +98,8 @@
 
             class Test {
                 void test() {
-                    if (org.example.ProtoLogImpl.Cache.TEST_GROUP_enabled[3]) { long protoLogParam0 = 100; double protoLogParam1 = 0.1; org.example.ProtoLogImpl.w(TEST_GROUP, -1473209266730422156L, 9, "test %d %f", protoLogParam0, protoLogParam1); } /* ProtoLog.w(TEST_GROUP, "test %d %f", 100, 0.1); */ if (org.example.ProtoLogImpl.Cache.TEST_GROUP_enabled[3]) { long protoLogParam0 = 100; double protoLogParam1 = 0.1; org.example.ProtoLogImpl.w(TEST_GROUP, -1473209266730422156L, 9, "test %d %f", protoLogParam0, protoLogParam1); }
-                    if (org.example.ProtoLogImpl.Cache.TEST_GROUP_enabled[3]) { long protoLogParam0 = 100; double protoLogParam1 = 0.1; org.example.ProtoLogImpl.w(TEST_GROUP, -1473209266730422156L, 9, "test %d %f", protoLogParam0, protoLogParam1); }
+                    if (org.example.ProtoLogImpl.Cache.TEST_GROUP_enabled[3]) { long protoLogParam0 = 100; double protoLogParam1 = 0.1; org.example.ProtoLogImpl.w(TEST_GROUP, -1473209266730422156L, 9, protoLogParam0, protoLogParam1); } /* ProtoLog.w(TEST_GROUP, "test %d %f", 100, 0.1); */ if (org.example.ProtoLogImpl.Cache.TEST_GROUP_enabled[3]) { long protoLogParam0 = 100; double protoLogParam1 = 0.1; org.example.ProtoLogImpl.w(TEST_GROUP, -1473209266730422156L, 9, protoLogParam0, protoLogParam1); }
+                    if (org.example.ProtoLogImpl.Cache.TEST_GROUP_enabled[3]) { long protoLogParam0 = 100; double protoLogParam1 = 0.1; org.example.ProtoLogImpl.w(TEST_GROUP, -1473209266730422156L, 9, protoLogParam0, protoLogParam1); }
                 }
             }
             """.trimIndent()
@@ -109,7 +109,7 @@
 
             class Test {
                 void test() {
-                    if (org.example.ProtoLogImpl.Cache.TEST_GROUP_enabled[3]) { org.example.ProtoLogImpl.w(TEST_GROUP, 3218600869538902408L, 0, "test", (Object[]) null); }
+                    if (org.example.ProtoLogImpl.Cache.TEST_GROUP_enabled[3]) { org.example.ProtoLogImpl.w(TEST_GROUP, 3218600869538902408L, 0, (Object[]) null); }
                 }
             }
             """.trimIndent()
@@ -119,7 +119,7 @@
 
             class Test {
                 void test() {
-                    if (org.example.ProtoLogImpl.Cache.TEST_GROUP_enabled[3]) { long protoLogParam0 = 100; double protoLogParam1 = 0.1; org.example.ProtoLogImpl.w(TEST_GROUP, -1473209266730422156L, 9, null, protoLogParam0, protoLogParam1); }
+                    if (org.example.ProtoLogImpl.Cache.TEST_GROUP_enabled[3]) { long protoLogParam0 = 100; double protoLogParam1 = 0.1; org.example.ProtoLogImpl.w(TEST_GROUP, -1473209266730422156L, 9, protoLogParam0, protoLogParam1); }
                 }
             }
             """.trimIndent()
@@ -129,7 +129,7 @@
 
             class Test {
                 void test() {
-                    if (org.example.ProtoLogImpl.Cache.TEST_GROUP_enabled[3]) { long protoLogParam0 = 100; double protoLogParam1 = 0.1; String protoLogParam2 = String.valueOf("test"); org.example.ProtoLogImpl.w(TEST_GROUP, -4447034859795564700L, 9, null, protoLogParam0, protoLogParam1, protoLogParam2); 
+                    if (org.example.ProtoLogImpl.Cache.TEST_GROUP_enabled[3]) { long protoLogParam0 = 100; double protoLogParam1 = 0.1; String protoLogParam2 = String.valueOf("test"); org.example.ProtoLogImpl.w(TEST_GROUP, -4447034859795564700L, 9, protoLogParam0, protoLogParam1, protoLogParam2); 
             
             }
                 }
@@ -172,13 +172,12 @@
         Truth.assertThat(protoLogCalls).hasSize(1)
         val methodCall = protoLogCalls[0] as MethodCallExpr
         assertEquals("w", methodCall.name.asString())
-        assertEquals(6, methodCall.arguments.size)
+        assertEquals(5, methodCall.arguments.size)
         assertEquals("TEST_GROUP", methodCall.arguments[0].toString())
         assertEquals("-1473209266730422156L", methodCall.arguments[1].toString())
         assertEquals(0b1001.toString(), methodCall.arguments[2].toString())
-        assertEquals("\"test %d %f\"", methodCall.arguments[3].toString())
-        assertEquals("protoLogParam0", methodCall.arguments[4].toString())
-        assertEquals("protoLogParam1", methodCall.arguments[5].toString())
+        assertEquals("protoLogParam0", methodCall.arguments[3].toString())
+        assertEquals("protoLogParam1", methodCall.arguments[4].toString())
         assertEquals(TRANSFORMED_CODE_TEXT_ENABLED, out)
     }
 
@@ -214,13 +213,12 @@
         Truth.assertThat(protoLogCalls).hasSize(3)
         val methodCall = protoLogCalls[0] as MethodCallExpr
         assertEquals("w", methodCall.name.asString())
-        assertEquals(6, methodCall.arguments.size)
+        assertEquals(5, methodCall.arguments.size)
         assertEquals("TEST_GROUP", methodCall.arguments[0].toString())
         assertEquals("-1473209266730422156L", methodCall.arguments[1].toString())
         assertEquals(0b1001.toString(), methodCall.arguments[2].toString())
-        assertEquals("\"test %d %f\"", methodCall.arguments[3].toString())
-        assertEquals("protoLogParam0", methodCall.arguments[4].toString())
-        assertEquals("protoLogParam1", methodCall.arguments[5].toString())
+        assertEquals("protoLogParam0", methodCall.arguments[3].toString())
+        assertEquals("protoLogParam1", methodCall.arguments[4].toString())
         assertEquals(TRANSFORMED_CODE_MULTICALL_TEXT, out)
     }
 
@@ -252,13 +250,13 @@
         Truth.assertThat(protoLogCalls).hasSize(1)
         val methodCall = protoLogCalls[0] as MethodCallExpr
         assertEquals("w", methodCall.name.asString())
-        assertEquals(7, methodCall.arguments.size)
+        assertEquals(6, methodCall.arguments.size)
         assertEquals("TEST_GROUP", methodCall.arguments[0].toString())
         assertEquals("-4447034859795564700L", methodCall.arguments[1].toString())
         assertEquals(0b001001.toString(), methodCall.arguments[2].toString())
-        assertEquals("protoLogParam0", methodCall.arguments[4].toString())
-        assertEquals("protoLogParam1", methodCall.arguments[5].toString())
-        assertEquals("protoLogParam2", methodCall.arguments[6].toString())
+        assertEquals("protoLogParam0", methodCall.arguments[3].toString())
+        assertEquals("protoLogParam1", methodCall.arguments[4].toString())
+        assertEquals("protoLogParam2", methodCall.arguments[5].toString())
         assertEquals(TRANSFORMED_CODE_MULTILINE_TEXT_ENABLED, out)
     }
 
@@ -289,7 +287,7 @@
         Truth.assertThat(protoLogCalls).hasSize(1)
         val methodCall = protoLogCalls[0] as MethodCallExpr
         assertEquals("w", methodCall.name.asString())
-        assertEquals(5, methodCall.arguments.size)
+        assertEquals(4, methodCall.arguments.size)
         assertEquals("TEST_GROUP", methodCall.arguments[0].toString())
         assertEquals("3218600869538902408L", methodCall.arguments[1].toString())
         assertEquals(0.toString(), methodCall.arguments[2].toString())
@@ -323,13 +321,12 @@
         Truth.assertThat(protoLogCalls).hasSize(1)
         val methodCall = protoLogCalls[0] as MethodCallExpr
         assertEquals("w", methodCall.name.asString())
-        assertEquals(6, methodCall.arguments.size)
+        assertEquals(5, methodCall.arguments.size)
         assertEquals("TEST_GROUP", methodCall.arguments[0].toString())
         assertEquals("-1473209266730422156L", methodCall.arguments[1].toString())
         assertEquals(0b1001.toString(), methodCall.arguments[2].toString())
-        assertEquals("null", methodCall.arguments[3].toString())
-        assertEquals("protoLogParam0", methodCall.arguments[4].toString())
-        assertEquals("protoLogParam1", methodCall.arguments[5].toString())
+        assertEquals("protoLogParam0", methodCall.arguments[3].toString())
+        assertEquals("protoLogParam1", methodCall.arguments[4].toString())
         assertEquals(TRANSFORMED_CODE_TEXT_DISABLED, out)
     }
 
@@ -361,14 +358,13 @@
         Truth.assertThat(protoLogCalls).hasSize(1)
         val methodCall = protoLogCalls[0] as MethodCallExpr
         assertEquals("w", methodCall.name.asString())
-        assertEquals(7, methodCall.arguments.size)
+        assertEquals(6, methodCall.arguments.size)
         assertEquals("TEST_GROUP", methodCall.arguments[0].toString())
         assertEquals("-4447034859795564700L", methodCall.arguments[1].toString())
         assertEquals(0b001001.toString(), methodCall.arguments[2].toString())
-        assertEquals("null", methodCall.arguments[3].toString())
-        assertEquals("protoLogParam0", methodCall.arguments[4].toString())
-        assertEquals("protoLogParam1", methodCall.arguments[5].toString())
-        assertEquals("protoLogParam2", methodCall.arguments[6].toString())
+        assertEquals("protoLogParam0", methodCall.arguments[3].toString())
+        assertEquals("protoLogParam1", methodCall.arguments[4].toString())
+        assertEquals("protoLogParam2", methodCall.arguments[5].toString())
         assertEquals(TRANSFORMED_CODE_MULTILINE_TEXT_DISABLED, out)
     }
 }