diff --git a/Android.bp b/Android.bp
index a3e39018..485f1be 100644
--- a/Android.bp
+++ b/Android.bp
@@ -149,6 +149,7 @@
         ":framework-javastream-protos",
         ":statslog-framework-java-gen", // FrameworkStatsLog.java
         ":audio_policy_configuration_V7_0",
+        ":perfetto_trace_javastream_protos",
     ],
 }
 
diff --git a/core/java/android/credentials/CredentialManager.java b/core/java/android/credentials/CredentialManager.java
index 47ee76e..796a57b 100644
--- a/core/java/android/credentials/CredentialManager.java
+++ b/core/java/android/credentials/CredentialManager.java
@@ -33,12 +33,12 @@
 import android.content.IntentSender;
 import android.os.Binder;
 import android.os.CancellationSignal;
+import android.os.IBinder;
 import android.os.ICancellationSignal;
 import android.os.OutcomeReceiver;
 import android.os.RemoteException;
 import android.provider.DeviceConfig;
 import android.util.Log;
-import android.view.autofill.IAutoFillManagerClient;
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
@@ -138,7 +138,7 @@
             @CallbackExecutor @NonNull Executor executor,
             @NonNull OutcomeReceiver<GetCandidateCredentialsResponse,
                     GetCandidateCredentialsException> callback,
-            @NonNull IAutoFillManagerClient clientCallback
+            @NonNull IBinder clientCallback
     ) {
         requireNonNull(request, "request must not be null");
         requireNonNull(callingPackage, "callingPackage must not be null");
diff --git a/core/java/android/credentials/ICredentialManager.aidl b/core/java/android/credentials/ICredentialManager.aidl
index 726bc97..d4a2d97 100644
--- a/core/java/android/credentials/ICredentialManager.aidl
+++ b/core/java/android/credentials/ICredentialManager.aidl
@@ -22,7 +22,6 @@
 import android.credentials.ClearCredentialStateRequest;
 import android.credentials.CreateCredentialRequest;
 import android.credentials.GetCandidateCredentialsRequest;
-import android.view.autofill.IAutoFillManagerClient;
 import android.credentials.GetCredentialRequest;
 import android.credentials.RegisterCredentialDescriptionRequest;
 import android.credentials.UnregisterCredentialDescriptionRequest;
@@ -33,6 +32,7 @@
 import android.credentials.IPrepareGetCredentialCallback;
 import android.credentials.ISetEnabledProvidersCallback;
 import android.content.ComponentName;
+import android.os.IBinder;
 import android.os.ICancellationSignal;
 
 /**
@@ -48,7 +48,7 @@
 
     @nullable ICancellationSignal executeCreateCredential(in CreateCredentialRequest request, in ICreateCredentialCallback callback, String callingPackage);
 
-    @nullable ICancellationSignal getCandidateCredentials(in GetCredentialRequest request, in IGetCandidateCredentialsCallback callback, in IAutoFillManagerClient clientCallback, String callingPackage);
+    @nullable ICancellationSignal getCandidateCredentials(in GetCredentialRequest request, in IGetCandidateCredentialsCallback callback, in IBinder clientCallback, String callingPackage);
 
     @nullable ICancellationSignal clearCredentialState(in ClearCredentialStateRequest request, in IClearCredentialStateCallback callback, String callingPackage);
 
diff --git a/core/java/android/provider/Settings.java b/core/java/android/provider/Settings.java
index f974ef4..7d84bb3 100644
--- a/core/java/android/provider/Settings.java
+++ b/core/java/android/provider/Settings.java
@@ -1931,9 +1931,7 @@
      * A matching Activity will only be found if
      * {@link NotificationManager#areAutomaticZenRulesUserManaged()} is true.
      * <p>
-     * Input: Intent's data URI set with an application name, using the "package" schema (like
-     * "package:com.my.app").
-     * Input: The id of the rule, provided in {@link #EXTRA_AUTOMATIC_ZEN_RULE_ID}.
+     * Input: The id of the rule, provided in the {@link #EXTRA_AUTOMATIC_ZEN_RULE_ID} extra.
      * <p>
      * Output: Nothing.
      */
diff --git a/core/java/android/tracing/perfetto/CreateIncrementalStateArgs.java b/core/java/android/tracing/perfetto/CreateIncrementalStateArgs.java
new file mode 100644
index 0000000..819cc8c
--- /dev/null
+++ b/core/java/android/tracing/perfetto/CreateIncrementalStateArgs.java
@@ -0,0 +1,43 @@
+/*
+ * 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 android.tracing.perfetto;
+
+import android.annotation.Nullable;
+
+/**
+ * @hide
+ * @param <DataSourceInstanceType> The type of datasource instance this state applied to.
+ */
+public class CreateIncrementalStateArgs<DataSourceInstanceType extends DataSourceInstance> {
+    private final DataSource<DataSourceInstanceType, Object, Object> mDataSource;
+    private final int mInstanceIndex;
+
+    CreateIncrementalStateArgs(DataSource dataSource, int instanceIndex) {
+        this.mDataSource = dataSource;
+        this.mInstanceIndex = instanceIndex;
+    }
+
+    /**
+     * Gets the datasource instance for this state with a lock.
+     * releaseDataSourceInstanceLocked must be called before this can be called again.
+     * @return The data source instance for this state.
+     *         Null if the datasource instance no longer exists.
+     */
+    public @Nullable DataSourceInstanceType getDataSourceInstanceLocked() {
+        return mDataSource.getDataSourceInstanceLocked(mInstanceIndex);
+    }
+}
diff --git a/core/java/android/tracing/perfetto/CreateTlsStateArgs.java b/core/java/android/tracing/perfetto/CreateTlsStateArgs.java
new file mode 100644
index 0000000..3fad2d1
--- /dev/null
+++ b/core/java/android/tracing/perfetto/CreateTlsStateArgs.java
@@ -0,0 +1,43 @@
+/*
+ * 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 android.tracing.perfetto;
+
+import android.annotation.Nullable;
+
+/**
+ * @hide
+ * @param <DataSourceInstanceType> The type of datasource instance this state applied to.
+ */
+public class CreateTlsStateArgs<DataSourceInstanceType extends DataSourceInstance> {
+    private final DataSource<DataSourceInstanceType, Object, Object> mDataSource;
+    private final int mInstanceIndex;
+
+    CreateTlsStateArgs(DataSource dataSource, int instanceIndex) {
+        this.mDataSource = dataSource;
+        this.mInstanceIndex = instanceIndex;
+    }
+
+    /**
+     * Gets the datasource instance for this state with a lock.
+     * releaseDataSourceInstanceLocked must be called before this can be called again.
+     * @return The data source instance for this state.
+     *         Null if the datasource instance no longer exists.
+     */
+    public @Nullable DataSourceInstanceType getDataSourceInstanceLocked() {
+        return mDataSource.getDataSourceInstanceLocked(mInstanceIndex);
+    }
+}
diff --git a/core/java/android/tracing/perfetto/DataSource.java b/core/java/android/tracing/perfetto/DataSource.java
new file mode 100644
index 0000000..4e08aee
--- /dev/null
+++ b/core/java/android/tracing/perfetto/DataSource.java
@@ -0,0 +1,163 @@
+/*
+ * 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 android.tracing.perfetto;
+
+import android.util.proto.ProtoInputStream;
+
+/**
+ * Templated base class meant to be derived by embedders to create a custom data
+ * source.
+ *
+ * @param <DataSourceInstanceType> The type for the DataSource instances that will be created from
+ *                                 this DataSource type.
+ * @param <TlsStateType> The type of the custom TLS state, if any is used.
+ * @param <IncrementalStateType> The type of the custom incremental state, if any is used.
+ *
+ * @hide
+ */
+public abstract class DataSource<DataSourceInstanceType extends DataSourceInstance,
+        TlsStateType, IncrementalStateType> {
+    protected final long mNativeObj;
+
+    public final String name;
+
+    /**
+     * A function implemented by each datasource to create a new data source instance.
+     *
+     * @param configStream A ProtoInputStream to read the tracing instance's config.
+     * @return A new data source instance setup with the provided config.
+     */
+    public abstract DataSourceInstanceType createInstance(
+            ProtoInputStream configStream, int instanceIndex);
+
+    /**
+     * Constructor for datasource base class.
+     *
+     * @param name The fully qualified name of the datasource.
+     */
+    public DataSource(String name) {
+        this.name = name;
+        this.mNativeObj = nativeCreate(this, name);
+    }
+
+    /**
+     * The main tracing method. Tracing code should call this passing a lambda as
+     * argument, with the following signature: void(TraceContext).
+     * <p>
+     * The lambda will be called synchronously (i.e., always before trace()
+     * returns) only if tracing is enabled and the data source has been enabled in
+     * the tracing config.
+     * <p>
+     * The lambda can be called more than once per trace() call, in the case of
+     * concurrent tracing sessions (or even if the data source is instantiated
+     * twice within the same trace config).
+     *
+     * @param fun The tracing lambda that will be called with the tracing contexts of each active
+     *            tracing instance.
+     */
+    public final void trace(
+            TraceFunction<DataSourceInstanceType, TlsStateType, IncrementalStateType> fun) {
+        nativeTrace(mNativeObj, fun);
+    }
+
+    /**
+     * Flush any trace data from this datasource that has not yet been flushed.
+     */
+    public final void flush() {
+        nativeFlushAll(mNativeObj);
+    }
+
+    /**
+     * Override this method to create a custom TlsState object for your DataSource. A new instance
+     * will be created per trace instance per thread.
+     *
+     * NOTE: Should only be called from native side.
+     */
+    protected TlsStateType createTlsState(CreateTlsStateArgs<DataSourceInstanceType> args) {
+        return null;
+    }
+
+    /**
+     * Override this method to create and use a custom IncrementalState object for your DataSource.
+     *
+     * NOTE: Should only be called from native side.
+     */
+    protected IncrementalStateType createIncrementalState(
+            CreateIncrementalStateArgs<DataSourceInstanceType> args) {
+        return null;
+    }
+
+    /**
+     * Registers the data source on all tracing backends, including ones that
+     * connect after the registration. Doing so enables the data source to receive
+     * Setup/Start/Stop notifications and makes the trace() method work when
+     * tracing is enabled and the data source is selected.
+     * <p>
+     * NOTE: Once registered, we cannot unregister the data source. Therefore, we should avoid
+     * creating and registering data source where not strictly required. This is a fundamental
+     * limitation of Perfetto itself.
+     *
+     * @param params Params to initialize the datasource with.
+     */
+    public void register(DataSourceParams params) {
+        nativeRegisterDataSource(this.mNativeObj, params.bufferExhaustedPolicy);
+    }
+
+    /**
+     * Gets the datasource instance with a specified index.
+     * IMPORTANT: releaseDataSourceInstance must be called after using the datasource instance.
+     * @param instanceIndex The index of the datasource to lock and get.
+     * @return The DataSourceInstance at index instanceIndex.
+     *         Null if the datasource instance at the requested index doesn't exist.
+     */
+    public DataSourceInstanceType getDataSourceInstanceLocked(int instanceIndex) {
+        return (DataSourceInstanceType) nativeGetPerfettoInstanceLocked(mNativeObj, instanceIndex);
+    }
+
+    /**
+     * Unlock the datasource at the specified index.
+     * @param instanceIndex The index of the datasource to unlock.
+     */
+    protected void releaseDataSourceInstance(int instanceIndex) {
+        nativeReleasePerfettoInstanceLocked(mNativeObj, instanceIndex);
+    }
+
+    /**
+     * Called from native side when a new tracing instance starts.
+     *
+     * @param rawConfig byte array of the PerfettoConfig encoded proto.
+     * @return A new Java DataSourceInstance object.
+     */
+    private DataSourceInstanceType createInstance(byte[] rawConfig, int instanceIndex) {
+        final ProtoInputStream inputStream = new ProtoInputStream(rawConfig);
+        return this.createInstance(inputStream, instanceIndex);
+    }
+
+    private static native void nativeRegisterDataSource(
+            long dataSourcePtr, int bufferExhaustedPolicy);
+
+    private static native long nativeCreate(DataSource thiz, String name);
+    private static native void nativeTrace(
+            long nativeDataSourcePointer, TraceFunction traceFunction);
+    private static native void nativeFlushAll(long nativeDataSourcePointer);
+    private static native long nativeGetFinalizer();
+
+    private static native DataSourceInstance nativeGetPerfettoInstanceLocked(
+            long dataSourcePtr, int dsInstanceIdx);
+    private static native void nativeReleasePerfettoInstanceLocked(
+            long dataSourcePtr, int dsInstanceIdx);
+}
diff --git a/core/java/android/tracing/perfetto/DataSourceInstance.java b/core/java/android/tracing/perfetto/DataSourceInstance.java
new file mode 100644
index 0000000..4994501
--- /dev/null
+++ b/core/java/android/tracing/perfetto/DataSourceInstance.java
@@ -0,0 +1,72 @@
+/*
+ * 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 android.tracing.perfetto;
+
+/**
+ * @hide
+ */
+public abstract class DataSourceInstance implements AutoCloseable {
+    private final DataSource mDataSource;
+    private final int mInstanceIndex;
+
+    public DataSourceInstance(DataSource dataSource, int instanceIndex) {
+        this.mDataSource = dataSource;
+        this.mInstanceIndex = instanceIndex;
+    }
+
+    /**
+     * Executed when the tracing instance starts running.
+     * <p>
+     * NOTE: This callback executes on the Perfetto internal thread and is blocking.
+     *       Anything that is run in this callback should execute quickly.
+     *
+     * @param args Start arguments.
+     */
+    protected void onStart(StartCallbackArguments args) {}
+
+    /**
+     * Executed when a flush is triggered.
+     * <p>
+     * NOTE: This callback executes on the Perfetto internal thread and is blocking.
+     *       Anything that is run in this callback should execute quickly.
+     * @param args Flush arguments.
+     */
+    protected void onFlush(FlushCallbackArguments args) {}
+
+    /**
+     * Executed when the tracing instance is stopped.
+     * <p>
+     * NOTE: This callback executes on the Perfetto internal thread and is blocking.
+     *       Anything that is run in this callback should execute quickly.
+     * @param args Stop arguments.
+     */
+    protected void onStop(StopCallbackArguments args) {}
+
+    @Override
+    public final void close() {
+        this.release();
+    }
+
+    /**
+     * Release the lock on the datasource once you are finished using it.
+     * Only required to be called when instance was retrieved with
+     * `DataSource#getDataSourceInstanceLocked`.
+     */
+    public final void release() {
+        mDataSource.releaseDataSourceInstance(mInstanceIndex);
+    }
+}
diff --git a/core/java/android/tracing/perfetto/DataSourceParams.java b/core/java/android/tracing/perfetto/DataSourceParams.java
new file mode 100644
index 0000000..6cd04e3
--- /dev/null
+++ b/core/java/android/tracing/perfetto/DataSourceParams.java
@@ -0,0 +1,57 @@
+/*
+ * 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 android.tracing.perfetto;
+
+import android.annotation.IntDef;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+
+/**
+ * DataSource Parameters
+ *
+ * @hide
+ */
+public class DataSourceParams {
+    /**
+     * @hide
+     */
+    @IntDef(value = {
+        PERFETTO_DS_BUFFER_EXHAUSTED_POLICY_DROP,
+        PERFETTO_DS_BUFFER_EXHAUSTED_POLICY_STALL_AND_ABORT,
+    })
+    @Retention(RetentionPolicy.SOURCE)
+    public @interface PerfettoDsBufferExhausted {}
+
+    // If the data source runs out of space when trying to acquire a new chunk,
+    // it will drop data.
+    public static final int PERFETTO_DS_BUFFER_EXHAUSTED_POLICY_DROP = 0;
+
+    // If the data source runs out of space when trying to acquire a new chunk,
+    // it will stall, retry and eventually abort if a free chunk is not acquired
+    // after a while.
+    public static final int PERFETTO_DS_BUFFER_EXHAUSTED_POLICY_STALL_AND_ABORT = 1;
+
+    public static DataSourceParams DEFAULTS =
+            new DataSourceParams(PERFETTO_DS_BUFFER_EXHAUSTED_POLICY_DROP);
+
+    public DataSourceParams(@PerfettoDsBufferExhausted int bufferExhaustedPolicy) {
+        this.bufferExhaustedPolicy = bufferExhaustedPolicy;
+    }
+
+    public final @PerfettoDsBufferExhausted int bufferExhaustedPolicy;
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt b/core/java/android/tracing/perfetto/FlushCallbackArguments.java
similarity index 74%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt
copy to core/java/android/tracing/perfetto/FlushCallbackArguments.java
index f9cdc1b..ecf6aee 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt
+++ b/core/java/android/tracing/perfetto/FlushCallbackArguments.java
@@ -14,9 +14,10 @@
  * limitations under the License.
  */
 
-package com.android.systemui.scene.shared.model
+package android.tracing.perfetto;
 
-import com.android.systemui.kosmos.Kosmos
-
-val Kosmos.sceneContainerConfig by
-    Kosmos.Fixture { FakeSceneContainerConfigModule().sceneContainerConfig }
+/**
+ * @hide
+ */
+public class FlushCallbackArguments {
+}
diff --git a/core/java/android/tracing/perfetto/InitArguments.java b/core/java/android/tracing/perfetto/InitArguments.java
new file mode 100644
index 0000000..da8c273
--- /dev/null
+++ b/core/java/android/tracing/perfetto/InitArguments.java
@@ -0,0 +1,54 @@
+/*
+ * 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 android.tracing.perfetto;
+
+import android.annotation.IntDef;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+
+/**
+ * @hide
+ */
+public class InitArguments {
+    public final @PerfettoBackend int backends;
+
+    /**
+     * @hide
+     */
+    @IntDef(value = {
+            PERFETTO_BACKEND_IN_PROCESS,
+            PERFETTO_BACKEND_SYSTEM,
+    })
+    @Retention(RetentionPolicy.SOURCE)
+    public @interface PerfettoBackend {}
+
+    // The in-process tracing backend. Keeps trace buffers in the process memory.
+    public static final int PERFETTO_BACKEND_IN_PROCESS = (1 << 0);
+
+    // The system tracing backend. Connects to the system tracing service (e.g.
+    // on Linux/Android/Mac uses a named UNIX socket).
+    public static final int PERFETTO_BACKEND_SYSTEM = (1 << 1);
+
+    public static InitArguments DEFAULTS = new InitArguments(PERFETTO_BACKEND_SYSTEM);
+
+    public static InitArguments TESTING = new InitArguments(PERFETTO_BACKEND_IN_PROCESS);
+
+    public InitArguments(@PerfettoBackend int backends) {
+        this.backends = backends;
+    }
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt b/core/java/android/tracing/perfetto/Producer.java
similarity index 60%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt
copy to core/java/android/tracing/perfetto/Producer.java
index f9cdc1b..a1b3eb7 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt
+++ b/core/java/android/tracing/perfetto/Producer.java
@@ -14,9 +14,21 @@
  * limitations under the License.
  */
 
-package com.android.systemui.scene.shared.model
+package android.tracing.perfetto;
 
-import com.android.systemui.kosmos.Kosmos
+/**
+ * @hide
+ */
+public class Producer {
 
-val Kosmos.sceneContainerConfig by
-    Kosmos.Fixture { FakeSceneContainerConfigModule().sceneContainerConfig }
+    /**
+     * Initializes the global Perfetto producer.
+     *
+     * @param args arguments on how to initialize the Perfetto producer.
+     */
+    public static void init(InitArguments args) {
+        nativePerfettoProducerInit(args.backends);
+    }
+
+    private static native void nativePerfettoProducerInit(int backends);
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt b/core/java/android/tracing/perfetto/StartCallbackArguments.java
similarity index 74%
rename from packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt
rename to core/java/android/tracing/perfetto/StartCallbackArguments.java
index f9cdc1b..9739d27 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt
+++ b/core/java/android/tracing/perfetto/StartCallbackArguments.java
@@ -14,9 +14,10 @@
  * limitations under the License.
  */
 
-package com.android.systemui.scene.shared.model
+package android.tracing.perfetto;
 
-import com.android.systemui.kosmos.Kosmos
-
-val Kosmos.sceneContainerConfig by
-    Kosmos.Fixture { FakeSceneContainerConfigModule().sceneContainerConfig }
+/**
+ * @hide
+ */
+public class StartCallbackArguments {
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt b/core/java/android/tracing/perfetto/StopCallbackArguments.java
similarity index 74%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt
copy to core/java/android/tracing/perfetto/StopCallbackArguments.java
index f9cdc1b..0cd1a18 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt
+++ b/core/java/android/tracing/perfetto/StopCallbackArguments.java
@@ -14,9 +14,10 @@
  * limitations under the License.
  */
 
-package com.android.systemui.scene.shared.model
+package android.tracing.perfetto;
 
-import com.android.systemui.kosmos.Kosmos
-
-val Kosmos.sceneContainerConfig by
-    Kosmos.Fixture { FakeSceneContainerConfigModule().sceneContainerConfig }
+/**
+ * @hide
+ */
+public class StopCallbackArguments {
+}
diff --git a/core/java/android/tracing/perfetto/TraceFunction.java b/core/java/android/tracing/perfetto/TraceFunction.java
new file mode 100644
index 0000000..62941df
--- /dev/null
+++ b/core/java/android/tracing/perfetto/TraceFunction.java
@@ -0,0 +1,43 @@
+/*
+ * 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 android.tracing.perfetto;
+
+import java.io.IOException;
+
+/**
+ * The interface for the trace function called from native on a trace call with a context.
+ *
+ * @param <DataSourceInstanceType> The type of DataSource this tracing context is for.
+ * @param <TlsStateType> The type of the custom TLS state, if any is used.
+ * @param <IncrementalStateType> The type of the custom incremental state, if any is used.
+ *
+ * @hide
+ */
+public interface TraceFunction<DataSourceInstanceType extends DataSourceInstance,
+        TlsStateType, IncrementalStateType> {
+
+    /**
+     * This function will be called synchronously (i.e., always before trace() returns) only if
+     * tracing is enabled and the data source has been enabled in the tracing config.
+     * It can be called more than once per trace() call, in the case of concurrent tracing sessions
+     * (or even if the data source is instantiated twice within the same trace config).
+     *
+     * @param ctx the tracing context to trace for in the trace function.
+     */
+    void trace(TracingContext<DataSourceInstanceType, TlsStateType, IncrementalStateType> ctx)
+            throws IOException;
+}
diff --git a/core/java/android/tracing/perfetto/TracingContext.java b/core/java/android/tracing/perfetto/TracingContext.java
new file mode 100644
index 0000000..0bce26e
--- /dev/null
+++ b/core/java/android/tracing/perfetto/TracingContext.java
@@ -0,0 +1,110 @@
+/*
+ * 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 android.tracing.perfetto;
+
+import android.util.proto.ProtoOutputStream;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Argument passed to the lambda function passed to Trace().
+ *
+ * @param <DataSourceInstanceType> The type of the datasource this tracing context is for.
+ * @param <TlsStateType> The type of the custom TLS state, if any is used.
+ * @param <IncrementalStateType> The type of the custom incremental state, if any is used.
+ *
+ * @hide
+ */
+public class TracingContext<DataSourceInstanceType extends DataSourceInstance,
+        TlsStateType, IncrementalStateType> {
+
+    private final long mContextPtr;
+    private final TlsStateType mTlsState;
+    private final IncrementalStateType mIncrementalState;
+    private final List<ProtoOutputStream> mTracePackets = new ArrayList<>();
+
+    // Should only be created from the native side.
+    private TracingContext(long contextPtr, TlsStateType tlsState,
+            IncrementalStateType incrementalState) {
+        this.mContextPtr = contextPtr;
+        this.mTlsState = tlsState;
+        this.mIncrementalState = incrementalState;
+    }
+
+    /**
+     * Creates a new output stream to be used to write a trace packet to. The output stream will be
+     * encoded to the proto binary representation when the callback trace function finishes and
+     * send over to the native side to be included in the proto buffer.
+     *
+     * @return A proto output stream to write a trace packet proto to
+     */
+    public ProtoOutputStream newTracePacket() {
+        final ProtoOutputStream os = new ProtoOutputStream(0);
+        mTracePackets.add(os);
+        return os;
+    }
+
+    /**
+     * Forces a commit of the thread-local tracing data written so far to the
+     * service. This is almost never required (tracing data is periodically
+     * committed as trace pages are filled up) and has a non-negligible
+     * performance hit (requires an IPC + refresh of the current thread-local
+     * chunk). The only case when this should be used is when handling OnStop()
+     * asynchronously, to ensure sure that the data is committed before the
+     * Stop timeout expires.
+     */
+    public void flush() {
+        nativeFlush(this, mContextPtr);
+    }
+
+    /**
+     * Can optionally be used to store custom per-sequence
+     * session data, which is not reset when incremental state is cleared
+     * (e.g. configuration options).
+     *
+     * @return The TlsState instance for the tracing thread and instance.
+     */
+    public TlsStateType getCustomTlsState() {
+        return this.mTlsState;
+    }
+
+    /**
+     * Can optionally be used store custom per-sequence
+     * incremental data (e.g., interning tables).
+     *
+     * @return The current IncrementalState object instance.
+     */
+    public IncrementalStateType getIncrementalState() {
+        return this.mIncrementalState;
+    }
+
+    // Called from native to get trace packets
+    private byte[][] getAndClearAllPendingTracePackets() {
+        byte[][] res = new byte[mTracePackets.size()][];
+        for (int i = 0; i < mTracePackets.size(); i++) {
+            ProtoOutputStream tracePacket = mTracePackets.get(i);
+            res[i] = tracePacket.getBytes();
+        }
+
+        mTracePackets.clear();
+        return res;
+    }
+
+    // private static native void nativeFlush(long nativeDataSourcePointer);
+    private static native void nativeFlush(TracingContext thiz, long ctxPointer);
+}
diff --git a/core/java/android/webkit/URLUtil.java b/core/java/android/webkit/URLUtil.java
index 828ec26..c6271d2 100644
--- a/core/java/android/webkit/URLUtil.java
+++ b/core/java/android/webkit/URLUtil.java
@@ -16,20 +16,40 @@
 
 package android.webkit;
 
+import android.annotation.FlaggedApi;
+import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.compat.Compatibility;
+import android.compat.annotation.ChangeId;
+import android.compat.annotation.EnabledSince;
 import android.compat.annotation.UnsupportedAppUsage;
 import android.net.ParseException;
 import android.net.Uri;
 import android.net.WebAddress;
+import android.os.Build;
 import android.util.Log;
 
 import java.io.UnsupportedEncodingException;
+import java.net.URLDecoder;
+import java.nio.charset.Charset;
 import java.util.Locale;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
 public final class URLUtil {
 
+    /**
+     * This feature enables parsing of Content-Disposition headers that conform to RFC 6266. In
+     * particular, this enables parsing of {@code filename*} values which can use a different
+     * character encoding.
+     *
+     * @hide
+     */
+    @ChangeId
+    @EnabledSince(targetSdkVersion = Build.VERSION_CODES.VANILLA_ICE_CREAM)
+    @FlaggedApi(android.os.Flags.FLAG_ANDROID_OS_BUILD_VANILLA_ICE_CREAM)
+    public static final long PARSE_CONTENT_DISPOSITION_USING_RFC_6266 = 319400769L;
+
     private static final String LOGTAG = "webkit";
     private static final boolean TRACE = false;
 
@@ -293,21 +313,58 @@
 
     /**
      * Guesses canonical filename that a download would have, using the URL and contentDisposition.
-     * File extension, if not defined, is added based on the mimetype
+     *
+     * <p>File extension, if not defined, is added based on the mimetype.
+     *
+     * <p>The {@code contentDisposition} argument will be treated differently depending on
+     * targetSdkVersion.
+     *
+     * <ul>
+     *   <li>For targetSDK versions &lt; {@code VANILLA_ICE_CREAM} it will be parsed based on RFC
+     *       2616.
+     *   <li>For targetSDK versions &gt;= {@code VANILLA_ICE_CREAM} it will be parsed based on RFC
+     *       6266.
+     * </ul>
+     *
+     * In practice, this means that from {@code VANILLA_ICE_CREAM}, this method will be able to
+     * parse {@code filename*} directives in the {@code contentDisposition} string.
+     *
+     * <p>The function also changed in the following ways in {@code VANILLA_ICE_CREAM}:
+     *
+     * <ul>
+     *   <li>If the suggested file type extension doesn't match the passed {@code mimeType}, the
+     *       method will append the appropriate extension instead of replacing the current
+     *       extension.
+     *   <li>If the suggested file name contains a path separator ({@code "/"}), the method will
+     *       replace this with the underscore character ({@code "_"}) instead of splitting the
+     *       result and only using the last part.
+     * </ul>
      *
      * @param url Url to the content
      * @param contentDisposition Content-Disposition HTTP header or {@code null}
      * @param mimeType Mime-type of the content or {@code null}
      * @return suggested filename
      */
-    public static final String guessFileName(
+    public static String guessFileName(
+            String url, @Nullable String contentDisposition, @Nullable String mimeType) {
+        if (android.os.Flags.androidOsBuildVanillaIceCream()) {
+            if (Compatibility.isChangeEnabled(PARSE_CONTENT_DISPOSITION_USING_RFC_6266)) {
+                return guessFileNameRfc6266(url, contentDisposition, mimeType);
+            }
+        }
+
+        return guessFileNameRfc2616(url, contentDisposition, mimeType);
+    }
+
+    /** Legacy implementation of guessFileName, based on RFC 2616. */
+    private static String guessFileNameRfc2616(
             String url, @Nullable String contentDisposition, @Nullable String mimeType) {
         String filename = null;
         String extension = null;
 
         // If we couldn't do anything with the hint, move toward the content disposition
         if (contentDisposition != null) {
-            filename = parseContentDisposition(contentDisposition);
+            filename = parseContentDispositionRfc2616(contentDisposition);
             if (filename != null) {
                 int index = filename.lastIndexOf('/') + 1;
                 if (index > 0) {
@@ -384,6 +441,128 @@
         return filename + extension;
     }
 
+    /**
+     * Guesses canonical filename that a download would have, using the URL and contentDisposition.
+     * Uses RFC 6266 for parsing the contentDisposition header value.
+     */
+    @NonNull
+    private static String guessFileNameRfc6266(
+            @NonNull String url, @Nullable String contentDisposition, @Nullable String mimeType) {
+        String filename = getFilenameSuggestion(url, contentDisposition);
+        // Split filename between base and extension
+        // Add an extension if filename does not have one
+        String extensionFromMimeType = suggestExtensionFromMimeType(mimeType);
+
+        if (filename.indexOf('.') < 0) {
+            // Filename does not have an extension, use the suggested one.
+            return filename + extensionFromMimeType;
+        }
+
+        // Filename already contains at least one dot.
+        // Compare the last segment of the extension against the mime type.
+        // If there's a mismatch, add the suggested extension instead.
+        if (mimeType != null && extensionDifferentFromMimeType(filename, mimeType)) {
+            return filename + extensionFromMimeType;
+        }
+        return filename;
+    }
+
+    /**
+     * Get the suggested file name from the {@code contentDisposition} or {@code url}. Will ensure
+     * that the filename contains no path separators by replacing them with the {@code "_"}
+     * character.
+     */
+    @NonNull
+    private static String getFilenameSuggestion(String url, @Nullable String contentDisposition) {
+        // First attempt to parse the Content-Disposition header if available
+        if (contentDisposition != null) {
+            String filename = getFilenameFromContentDispositionRfc6266(contentDisposition);
+            if (filename != null) {
+                return replacePathSeparators(filename);
+            }
+        }
+
+        // Try to generate a filename based on the URL.
+        if (url != null) {
+            Uri parsedUri = Uri.parse(url);
+            String lastPathSegment = parsedUri.getLastPathSegment();
+            if (lastPathSegment != null) {
+                return replacePathSeparators(lastPathSegment);
+            }
+        }
+
+        // Finally, if couldn't get filename from URI, get a generic filename.
+        return "downloadfile";
+    }
+
+    /**
+     * Replace all instances of {@code "/"} with {@code "_"} to avoid filenames that navigate the
+     * path.
+     */
+    @NonNull
+    private static String replacePathSeparators(@NonNull String raw) {
+        return raw.replaceAll("/", "_");
+    }
+
+    /**
+     * Check if the {@code filename} has an extension that is different from the expected one based
+     * on the {@code mimeType}.
+     */
+    private static boolean extensionDifferentFromMimeType(
+            @NonNull String filename, @NonNull String mimeType) {
+        int lastDotIndex = filename.lastIndexOf('.');
+        String typeFromExt =
+                MimeTypeMap.getSingleton()
+                        .getMimeTypeFromExtension(filename.substring(lastDotIndex + 1));
+        return typeFromExt != null && !typeFromExt.equalsIgnoreCase(mimeType);
+    }
+
+    /**
+     * Get a candidate file extension (including the {@code .}) for the given mimeType. will return
+     * {@code ".bin"} if {@code mimeType} is {@code null}
+     *
+     * @param mimeType Reported mimetype
+     * @return A file extension, including the {@code .}
+     */
+    @NonNull
+    private static String suggestExtensionFromMimeType(@Nullable String mimeType) {
+        if (mimeType == null) {
+            return ".bin";
+        }
+        String extensionFromMimeType =
+                MimeTypeMap.getSingleton().getExtensionFromMimeType(mimeType);
+        if (extensionFromMimeType != null) {
+            return "." + extensionFromMimeType;
+        }
+        if (mimeType.equalsIgnoreCase("text/html")) {
+            return ".html";
+        } else if (mimeType.toLowerCase(Locale.ROOT).startsWith("text/")) {
+            return ".txt";
+        } else {
+            return ".bin";
+        }
+    }
+
+    /**
+     * Parse the Content-Disposition HTTP Header.
+     *
+     * <p>Behavior depends on targetSdkVersion.
+     *
+     * <ul>
+     *   <li>For targetSDK versions &lt; {@code VANILLA_ICE_CREAM} it will parse based on RFC 2616.
+     *   <li>For targetSDK versions &gt;= {@code VANILLA_ICE_CREAM} it will parse based on RFC 6266.
+     * </ul>
+     */
+    @UnsupportedAppUsage
+    static String parseContentDisposition(String contentDisposition) {
+        if (android.os.Flags.androidOsBuildVanillaIceCream()) {
+            if (Compatibility.isChangeEnabled(PARSE_CONTENT_DISPOSITION_USING_RFC_6266)) {
+                return getFilenameFromContentDispositionRfc6266(contentDisposition);
+            }
+        }
+        return parseContentDispositionRfc2616(contentDisposition);
+    }
+
     /** Regex used to parse content-disposition headers */
     private static final Pattern CONTENT_DISPOSITION_PATTERN =
             Pattern.compile(
@@ -391,15 +570,14 @@
                     Pattern.CASE_INSENSITIVE);
 
     /**
-     * Parse the Content-Disposition HTTP Header. The format of the header is defined here:
-     * http://www.w3.org/Protocols/rfc2616/rfc2616-sec19.html This header provides a filename for
-     * content that is going to be downloaded to the file system. We only support the attachment
-     * type. Note that RFC 2616 specifies the filename value must be double-quoted. Unfortunately
-     * some servers do not quote the value so to maintain consistent behaviour with other browsers,
-     * we allow unquoted values too.
+     * Parse the Content-Disposition HTTP Header. The format of the header is defined here: <a
+     * href="http://www.w3.org/Protocols/rfc2616/rfc2616-sec19.html">rfc2616 Section 19</a>. This
+     * header provides a filename for content that is going to be downloaded to the file system. We
+     * only support the attachment type. Note that RFC 2616 specifies the filename value must be
+     * double-quoted. Unfortunately some servers do not quote the value so to maintain consistent
+     * behaviour with other browsers, we allow unquoted values too.
      */
-    @UnsupportedAppUsage
-    static String parseContentDisposition(String contentDisposition) {
+    private static String parseContentDispositionRfc2616(String contentDisposition) {
         try {
             Matcher m = CONTENT_DISPOSITION_PATTERN.matcher(contentDisposition);
             if (m.find()) {
@@ -410,4 +588,136 @@
         }
         return null;
     }
+
+    /**
+     * Pattern for parsing individual content disposition key-value pairs.
+     *
+     * <p>The pattern will attempt to parse the value as either single-, double-, or unquoted. For
+     * the single- and double-quoted options, the pattern allows escaped quotes as part of the
+     * value, as per <a href="https://datatracker.ietf.org/doc/html/rfc2616#section-2.2">rfc2616
+     * section-2.2</a>
+     */
+    @SuppressWarnings("RegExpRepeatedSpace") // Spaces are only for readability.
+    private static final Pattern DISPOSITION_PATTERN =
+            Pattern.compile(
+                    """
+                            \\s*(\\S+?) # Group 1: parameter name
+                            \\s*=\\s* # Match equals sign
+                            (?: # non-capturing group of options
+                               '( (?: [^'\\\\] | \\\\. )* )' # Group 2: single-quoted
+                             | "( (?: [^"\\\\] | \\\\. )*  )" # Group 3: double-quoted
+                             | ( [^'"][^;\\s]* ) # Group 4: un-quoted parameter
+                            )\\s*;? # Optional end semicolon""",
+                    Pattern.COMMENTS);
+
+    /**
+     * Extract filename from a {@code Content-Disposition} header value.
+     *
+     * <p>This method implements the parsing defined in <a
+     * href="https://datatracker.ietf.org/doc/html/rfc6266">RFC 6266</a>, supporting both the {@code
+     * filename} and {@code filename*} disposition parameters. If the passed header value has the
+     * {@code "inline"} disposition type, this method will return {@code null} to indicate that a
+     * download was not intended.
+     *
+     * <p>If both {@code filename*} and {@code filename} is present, the former will be returned, as
+     * per the RFC. Invalid encoded values will be ignored.
+     *
+     * @param contentDisposition Value of {@code Content-Disposition} header.
+     * @return The filename suggested by the header or {@code null} if no filename could be parsed
+     *     from the header value.
+     */
+    @Nullable
+    private static String getFilenameFromContentDispositionRfc6266(
+            @NonNull String contentDisposition) {
+        String[] parts = contentDisposition.trim().split(";", 2);
+        if (parts.length < 2) {
+            // Need at least 2 parts, the `disposition-type` and at least one `disposition-parm`.
+            return null;
+        }
+        String dispositionType = parts[0].trim();
+        if ("inline".equalsIgnoreCase(dispositionType)) {
+            // "inline" should not result in a download.
+            // Unknown disposition types should be handles as "attachment"
+            // https://datatracker.ietf.org/doc/html/rfc6266#section-4.2
+            return null;
+        }
+        String dispositionParameters = parts[1];
+        Matcher matcher = DISPOSITION_PATTERN.matcher(dispositionParameters);
+        String filename = null;
+        String filenameExt = null;
+        while (matcher.find()) {
+            String parameter = matcher.group(1);
+            String value;
+            if (matcher.group(2) != null) {
+                value = removeSlashEscapes(matcher.group(2)); // Value was single-quoted
+            } else if (matcher.group(3) != null) {
+                value = removeSlashEscapes(matcher.group(3)); // Value was double-quoted
+            } else {
+                value = matcher.group(4); // Value was un-quoted
+            }
+
+            if (parameter == null || value == null) {
+                continue;
+            }
+
+            if ("filename*".equalsIgnoreCase(parameter)) {
+                filenameExt = parseExtValueString(value);
+            } else if ("filename".equalsIgnoreCase(parameter)) {
+                filename = value;
+            }
+        }
+
+        // RFC 6266 dictates the filenameExt should be preferred if present.
+        if (filenameExt != null) {
+            return filenameExt;
+        }
+        return filename;
+    }
+
+    /** Replace escapes of the \X form with X. */
+    private static String removeSlashEscapes(String raw) {
+        if (raw == null) {
+            return null;
+        }
+        return raw.replaceAll("\\\\(.)", "$1");
+    }
+
+    /**
+     * Parse an extended value string which can be percent-encoded. Return {@code} null if unable to
+     * parse the string.
+     */
+    private static String parseExtValueString(String raw) {
+        String[] parts = raw.split("'", 3);
+        if (parts.length < 3) {
+            return null;
+        }
+
+        String encoding = parts[0];
+        // Intentionally ignore parts[1] (language).
+        String valueChars = parts[2];
+
+        try {
+            // The URLDecoder force-decodes + as " "
+            // so preemptively replace all values with the encoded value to preserve them.
+            Charset charset = Charset.forName(encoding);
+            String valueWithEncodedPlus = encodePlusCharacters(valueChars, charset);
+            return URLDecoder.decode(valueWithEncodedPlus, charset);
+        } catch (RuntimeException ignored) {
+            return null; // Ignoring an un-parsable value is within spec.
+        }
+    }
+
+    /**
+     * Replace all instances of {@code "+"} with the percent-encoded equivalent for the given {@code
+     * charset}.
+     */
+    @NonNull
+    private static String encodePlusCharacters(@NonNull String valueChars, Charset charset) {
+        StringBuilder sb = new StringBuilder();
+        for (byte b : charset.encode("+").array()) {
+            // Formatting a byte is not possible with TextUtils.formatSimple
+            sb.append(String.format("%02x", b));
+        }
+        return valueChars.replaceAll("\\+", sb.toString());
+    }
 }
diff --git a/core/java/android/webkit/WebViewFactory.java b/core/java/android/webkit/WebViewFactory.java
index 0ef37d1..53b047a 100644
--- a/core/java/android/webkit/WebViewFactory.java
+++ b/core/java/android/webkit/WebViewFactory.java
@@ -16,6 +16,8 @@
 
 package android.webkit;
 
+import static android.webkit.Flags.updateServiceV2;
+
 import android.annotation.NonNull;
 import android.annotation.SystemApi;
 import android.annotation.UptimeMillisLong;
@@ -33,6 +35,7 @@
 import android.os.ServiceManager;
 import android.os.SystemClock;
 import android.os.Trace;
+import android.text.TextUtils;
 import android.util.AndroidRuntimeException;
 import android.util.ArraySet;
 import android.util.Log;
@@ -410,6 +413,21 @@
         }
     }
 
+    // Returns whether the given package is enabled.
+    // This state can be changed by the user from Settings->Apps
+    private static boolean isEnabledPackage(PackageInfo packageInfo) {
+        if (packageInfo == null) return false;
+        return packageInfo.applicationInfo.enabled;
+    }
+
+    // Return {@code true} if the package is installed and not hidden
+    private static boolean isInstalledPackage(PackageInfo packageInfo) {
+        if (packageInfo == null) return false;
+        return (((packageInfo.applicationInfo.flags & ApplicationInfo.FLAG_INSTALLED) != 0)
+                && ((packageInfo.applicationInfo.privateFlags & ApplicationInfo.PRIVATE_FLAG_HIDDEN)
+                        == 0));
+    }
+
     @UnsupportedAppUsage
     private static Context getWebViewContextAndSetProvider() throws MissingWebViewPackageException {
         Application initialApplication = AppGlobals.getInitialApplication();
@@ -456,6 +474,21 @@
                 Trace.traceEnd(Trace.TRACE_TAG_WEBVIEW);
             }
 
+            if (updateServiceV2() && !isInstalledPackage(newPackageInfo)) {
+                throw new MissingWebViewPackageException(
+                        TextUtils.formatSimple(
+                                "Current WebView Package (%s) is not installed for the current "
+                                + "user",
+                                newPackageInfo.packageName));
+            }
+
+            if (updateServiceV2() && !isEnabledPackage(newPackageInfo)) {
+                throw new MissingWebViewPackageException(
+                        TextUtils.formatSimple(
+                                "Current WebView Package (%s) is not enabled for the current user",
+                                newPackageInfo.packageName));
+            }
+
             // Validate the newly fetched package info, throws MissingWebViewPackageException on
             // failure
             verifyPackageInfo(response.packageInfo, newPackageInfo);
diff --git a/core/jni/Android.bp b/core/jni/Android.bp
index 2a744e3..c8fd246 100644
--- a/core/jni/Android.bp
+++ b/core/jni/Android.bp
@@ -269,6 +269,9 @@
                 "android_window_WindowInfosListener.cpp",
                 "android_window_ScreenCapture.cpp",
                 "jni_common.cpp",
+                "android_tracing_PerfettoDataSource.cpp",
+                "android_tracing_PerfettoDataSourceInstance.cpp",
+                "android_tracing_PerfettoProducer.cpp",
             ],
 
             static_libs: [
@@ -282,6 +285,7 @@
                 "libscrypt_static",
                 "libstatssocket_lazy",
                 "libskia",
+                "libperfetto_client_experimental",
             ],
 
             shared_libs: [
@@ -355,6 +359,7 @@
                 "server_configurable_flags",
                 "libimage_io",
                 "libultrahdr",
+                "libperfetto_c",
             ],
             export_shared_lib_headers: [
                 // our headers include libnativewindow's public headers
diff --git a/core/jni/AndroidRuntime.cpp b/core/jni/AndroidRuntime.cpp
index 17aad43..7a16318 100644
--- a/core/jni/AndroidRuntime.cpp
+++ b/core/jni/AndroidRuntime.cpp
@@ -220,6 +220,9 @@
 extern int register_android_window_WindowInfosListener(JNIEnv* env);
 extern int register_android_window_ScreenCapture(JNIEnv* env);
 extern int register_jni_common(JNIEnv* env);
+extern int register_android_tracing_PerfettoDataSource(JNIEnv* env);
+extern int register_android_tracing_PerfettoDataSourceInstance(JNIEnv* env);
+extern int register_android_tracing_PerfettoProducer(JNIEnv* env);
 
 // Namespace for Android Runtime flags applied during boot time.
 static const char* RUNTIME_NATIVE_BOOT_NAMESPACE = "runtime_native_boot";
@@ -1675,6 +1678,10 @@
         REG_JNI(register_android_window_WindowInfosListener),
         REG_JNI(register_android_window_ScreenCapture),
         REG_JNI(register_jni_common),
+
+        REG_JNI(register_android_tracing_PerfettoDataSource),
+        REG_JNI(register_android_tracing_PerfettoDataSourceInstance),
+        REG_JNI(register_android_tracing_PerfettoProducer),
 };
 
 /*
diff --git a/core/jni/android_tracing_PerfettoDataSource.cpp b/core/jni/android_tracing_PerfettoDataSource.cpp
new file mode 100644
index 0000000..d710698
--- /dev/null
+++ b/core/jni/android_tracing_PerfettoDataSource.cpp
@@ -0,0 +1,437 @@
+/*
+ * Copyright 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.
+ */
+
+#define LOG_TAG "Perfetto"
+
+#include "android_tracing_PerfettoDataSource.h"
+
+#include <android_runtime/AndroidRuntime.h>
+#include <android_runtime/Log.h>
+#include <nativehelper/JNIHelp.h>
+#include <perfetto/public/data_source.h>
+#include <perfetto/public/producer.h>
+#include <perfetto/public/protos/trace/test_event.pzc.h>
+#include <perfetto/public/protos/trace/trace_packet.pzc.h>
+#include <perfetto/tracing.h>
+#include <utils/Log.h>
+#include <utils/RefBase.h>
+
+#include <sstream>
+#include <thread>
+
+#include "core_jni_helpers.h"
+
+namespace android {
+
+static struct {
+    jclass clazz;
+    jmethodID createInstance;
+    jmethodID createTlsState;
+    jmethodID createIncrementalState;
+} gPerfettoDataSourceClassInfo;
+
+static struct {
+    jclass clazz;
+    jmethodID init;
+    jmethodID getAndClearAllPendingTracePackets;
+} gTracingContextClassInfo;
+
+static struct {
+    jclass clazz;
+    jmethodID init;
+} gCreateTlsStateArgsClassInfo;
+
+static struct {
+    jclass clazz;
+    jmethodID init;
+} gCreateIncrementalStateArgsClassInfo;
+
+static JavaVM* gVm;
+
+struct TlsState {
+    jobject jobj;
+};
+
+struct IncrementalState {
+    jobject jobj;
+};
+
+static void traceAllPendingPackets(JNIEnv* env, jobject jCtx, PerfettoDsTracerIterator* ctx) {
+    jobjectArray packets =
+            (jobjectArray)env
+                    ->CallObjectMethod(jCtx,
+                                       gTracingContextClassInfo.getAndClearAllPendingTracePackets);
+    if (env->ExceptionOccurred()) {
+        env->ExceptionDescribe();
+        env->ExceptionClear();
+
+        LOG_ALWAYS_FATAL("Failed to call java context finalize method");
+    }
+
+    int packets_count = env->GetArrayLength(packets);
+    for (int i = 0; i < packets_count; i++) {
+        jbyteArray packet_proto_buffer = (jbyteArray)env->GetObjectArrayElement(packets, i);
+
+        jbyte* raw_proto_buffer = env->GetByteArrayElements(packet_proto_buffer, 0);
+        int buffer_size = env->GetArrayLength(packet_proto_buffer);
+
+        struct PerfettoDsRootTracePacket trace_packet;
+        PerfettoDsTracerPacketBegin(ctx, &trace_packet);
+        PerfettoPbMsgAppendBytes(&trace_packet.msg.msg, (const uint8_t*)raw_proto_buffer,
+                                 buffer_size);
+        PerfettoDsTracerPacketEnd(ctx, &trace_packet);
+    }
+}
+
+PerfettoDataSource::PerfettoDataSource(JNIEnv* env, jobject javaDataSource,
+                                       std::string dataSourceName)
+      : dataSourceName(std::move(dataSourceName)),
+        mJavaDataSource(env->NewGlobalRef(javaDataSource)) {}
+
+jobject PerfettoDataSource::newInstance(JNIEnv* env, void* ds_config, size_t ds_config_size,
+                                        PerfettoDsInstanceIndex inst_id) {
+    jbyteArray configArray = env->NewByteArray(ds_config_size);
+
+    void* temp = env->GetPrimitiveArrayCritical((jarray)configArray, 0);
+    memcpy(temp, ds_config, ds_config_size);
+    env->ReleasePrimitiveArrayCritical(configArray, temp, 0);
+
+    jobject instance =
+            env->CallObjectMethod(mJavaDataSource, gPerfettoDataSourceClassInfo.createInstance,
+                                  configArray, inst_id);
+
+    if (env->ExceptionCheck()) {
+        LOGE_EX(env);
+        env->ExceptionClear();
+        LOG_ALWAYS_FATAL("Failed to create new Java Perfetto datasource instance");
+    }
+
+    return instance;
+}
+
+jobject PerfettoDataSource::createTlsStateGlobalRef(JNIEnv* env, PerfettoDsInstanceIndex inst_id) {
+    ScopedLocalRef<jobject> args(env,
+                                 env->NewObject(gCreateTlsStateArgsClassInfo.clazz,
+                                                gCreateTlsStateArgsClassInfo.init, mJavaDataSource,
+                                                inst_id));
+
+    ScopedLocalRef<jobject> tslState(env,
+                                     env->CallObjectMethod(mJavaDataSource,
+                                                           gPerfettoDataSourceClassInfo
+                                                                   .createTlsState,
+                                                           args.get()));
+
+    if (env->ExceptionCheck()) {
+        LOGE_EX(env);
+        env->ExceptionClear();
+        LOG_ALWAYS_FATAL("Failed to create new Java Perfetto incremental state");
+    }
+
+    return env->NewGlobalRef(tslState.get());
+}
+
+jobject PerfettoDataSource::createIncrementalStateGlobalRef(JNIEnv* env,
+                                                            PerfettoDsInstanceIndex inst_id) {
+    ScopedLocalRef<jobject> args(env,
+                                 env->NewObject(gCreateIncrementalStateArgsClassInfo.clazz,
+                                                gCreateIncrementalStateArgsClassInfo.init,
+                                                mJavaDataSource, inst_id));
+
+    ScopedLocalRef<jobject> incrementalState(env,
+                                             env->CallObjectMethod(mJavaDataSource,
+                                                                   gPerfettoDataSourceClassInfo
+                                                                           .createIncrementalState,
+                                                                   args.get()));
+
+    if (env->ExceptionCheck()) {
+        LOGE_EX(env);
+        env->ExceptionClear();
+        LOG_ALWAYS_FATAL("Failed to create Java Perfetto incremental state");
+    }
+
+    return env->NewGlobalRef(incrementalState.get());
+}
+
+void PerfettoDataSource::trace(JNIEnv* env, jobject traceFunction) {
+    PERFETTO_DS_TRACE(dataSource, ctx) {
+        TlsState* tls_state =
+                reinterpret_cast<TlsState*>(PerfettoDsGetCustomTls(&dataSource, &ctx));
+        IncrementalState* incr_state = reinterpret_cast<IncrementalState*>(
+                PerfettoDsGetIncrementalState(&dataSource, &ctx));
+
+        ScopedLocalRef<jobject> jCtx(env,
+                                     env->NewObject(gTracingContextClassInfo.clazz,
+                                                    gTracingContextClassInfo.init, &ctx,
+                                                    tls_state->jobj, incr_state->jobj));
+
+        jclass objclass = env->GetObjectClass(traceFunction);
+        jmethodID method =
+                env->GetMethodID(objclass, "trace", "(Landroid/tracing/perfetto/TracingContext;)V");
+        if (method == 0) {
+            LOG_ALWAYS_FATAL("Failed to get method id");
+        }
+
+        env->ExceptionClear();
+
+        env->CallVoidMethod(traceFunction, method, jCtx.get());
+        if (env->ExceptionOccurred()) {
+            env->ExceptionDescribe();
+            env->ExceptionClear();
+            LOG_ALWAYS_FATAL("Failed to call java trace method");
+        }
+
+        traceAllPendingPackets(env, jCtx.get(), &ctx);
+    }
+}
+
+void PerfettoDataSource::flushAll() {
+    PERFETTO_DS_TRACE(dataSource, ctx) {
+        PerfettoDsTracerFlush(&ctx, nullptr, nullptr);
+    }
+}
+
+PerfettoDataSource::~PerfettoDataSource() {
+    JNIEnv* env = AndroidRuntime::getJNIEnv();
+    env->DeleteWeakGlobalRef(mJavaDataSource);
+}
+
+jlong nativeCreate(JNIEnv* env, jclass clazz, jobject javaDataSource, jstring name) {
+    const char* nativeString = env->GetStringUTFChars(name, 0);
+    PerfettoDataSource* dataSource = new PerfettoDataSource(env, javaDataSource, nativeString);
+    env->ReleaseStringUTFChars(name, nativeString);
+
+    dataSource->incStrong((void*)nativeCreate);
+
+    return reinterpret_cast<jlong>(dataSource);
+}
+
+void nativeDestroy(void* ptr) {
+    PerfettoDataSource* dataSource = reinterpret_cast<PerfettoDataSource*>(ptr);
+    dataSource->decStrong((void*)nativeCreate);
+}
+
+static jlong nativeGetFinalizer(JNIEnv* /* env */, jclass /* clazz */) {
+    return static_cast<jlong>(reinterpret_cast<uintptr_t>(&nativeDestroy));
+}
+
+void nativeTrace(JNIEnv* env, jclass clazz, jlong dataSourcePtr, jobject traceFunctionInterface) {
+    sp<PerfettoDataSource> datasource = reinterpret_cast<PerfettoDataSource*>(dataSourcePtr);
+
+    datasource->trace(env, traceFunctionInterface);
+}
+
+void nativeFlush(JNIEnv* env, jclass clazz, jobject jCtx, jlong ctxPtr) {
+    auto* ctx = reinterpret_cast<struct PerfettoDsTracerIterator*>(ctxPtr);
+    traceAllPendingPackets(env, jCtx, ctx);
+    PerfettoDsTracerFlush(ctx, nullptr, nullptr);
+}
+
+void nativeFlushAll(JNIEnv* env, jclass clazz, jlong ptr) {
+    sp<PerfettoDataSource> datasource = reinterpret_cast<PerfettoDataSource*>(ptr);
+    datasource->flushAll();
+}
+
+void nativeRegisterDataSource(JNIEnv* env, jclass clazz, jlong datasource_ptr,
+                              int buffer_exhausted_policy) {
+    sp<PerfettoDataSource> datasource = reinterpret_cast<PerfettoDataSource*>(datasource_ptr);
+
+    struct PerfettoDsParams params = PerfettoDsParamsDefault();
+    params.buffer_exhausted_policy = (PerfettoDsBufferExhaustedPolicy)buffer_exhausted_policy;
+
+    params.user_arg = reinterpret_cast<void*>(datasource.get());
+
+    params.on_setup_cb = [](struct PerfettoDsImpl*, PerfettoDsInstanceIndex inst_id,
+                            void* ds_config, size_t ds_config_size, void* user_arg,
+                            struct PerfettoDsOnSetupArgs*) -> void* {
+        JNIEnv* env = GetOrAttachJNIEnvironment(gVm, JNI_VERSION_1_6);
+
+        auto* datasource = reinterpret_cast<PerfettoDataSource*>(user_arg);
+
+        ScopedLocalRef<jobject> java_data_source_instance(env,
+                                                          datasource->newInstance(env, ds_config,
+                                                                                  ds_config_size,
+                                                                                  inst_id));
+
+        auto* datasource_instance =
+                new PerfettoDataSourceInstance(env, java_data_source_instance.get(), inst_id);
+
+        return static_cast<void*>(datasource_instance);
+    };
+
+    params.on_create_tls_cb = [](struct PerfettoDsImpl* ds_impl, PerfettoDsInstanceIndex inst_id,
+                                 struct PerfettoDsTracerImpl* tracer, void* user_arg) -> void* {
+        JNIEnv* env = GetOrAttachJNIEnvironment(gVm, JNI_VERSION_1_6);
+
+        auto* datasource = reinterpret_cast<PerfettoDataSource*>(user_arg);
+
+        jobject java_tls_state = datasource->createTlsStateGlobalRef(env, inst_id);
+
+        auto* tls_state = new TlsState(java_tls_state);
+
+        return static_cast<void*>(tls_state);
+    };
+
+    params.on_delete_tls_cb = [](void* ptr) {
+        JNIEnv* env = GetOrAttachJNIEnvironment(gVm, JNI_VERSION_1_6);
+
+        TlsState* tls_state = reinterpret_cast<TlsState*>(ptr);
+        env->DeleteGlobalRef(tls_state->jobj);
+        delete tls_state;
+    };
+
+    params.on_create_incr_cb = [](struct PerfettoDsImpl* ds_impl, PerfettoDsInstanceIndex inst_id,
+                                  struct PerfettoDsTracerImpl* tracer, void* user_arg) -> void* {
+        JNIEnv* env = GetOrAttachJNIEnvironment(gVm, JNI_VERSION_1_6);
+
+        auto* datasource = reinterpret_cast<PerfettoDataSource*>(user_arg);
+        jobject java_incr_state = datasource->createIncrementalStateGlobalRef(env, inst_id);
+
+        auto* incr_state = new IncrementalState(java_incr_state);
+        return static_cast<void*>(incr_state);
+    };
+
+    params.on_delete_incr_cb = [](void* ptr) {
+        JNIEnv* env = GetOrAttachJNIEnvironment(gVm, JNI_VERSION_1_6);
+
+        IncrementalState* incr_state = reinterpret_cast<IncrementalState*>(ptr);
+        env->DeleteGlobalRef(incr_state->jobj);
+        delete incr_state;
+    };
+
+    params.on_start_cb = [](struct PerfettoDsImpl*, PerfettoDsInstanceIndex, void*, void* inst_ctx,
+                            struct PerfettoDsOnStartArgs*) {
+        JNIEnv* env = GetOrAttachJNIEnvironment(gVm, JNI_VERSION_1_6);
+
+        auto* datasource_instance = static_cast<PerfettoDataSourceInstance*>(inst_ctx);
+        datasource_instance->onStart(env);
+    };
+
+    params.on_flush_cb = [](struct PerfettoDsImpl*, PerfettoDsInstanceIndex, void*, void* inst_ctx,
+                            struct PerfettoDsOnFlushArgs*) {
+        JNIEnv* env = GetOrAttachJNIEnvironment(gVm, JNI_VERSION_1_6);
+
+        auto* datasource_instance = static_cast<PerfettoDataSourceInstance*>(inst_ctx);
+        datasource_instance->onFlush(env);
+    };
+
+    params.on_stop_cb = [](struct PerfettoDsImpl*, PerfettoDsInstanceIndex inst_id, void* user_arg,
+                           void* inst_ctx, struct PerfettoDsOnStopArgs*) {
+        JNIEnv* env = GetOrAttachJNIEnvironment(gVm, JNI_VERSION_1_6);
+
+        auto* datasource_instance = static_cast<PerfettoDataSourceInstance*>(inst_ctx);
+        datasource_instance->onStop(env);
+    };
+
+    params.on_destroy_cb = [](struct PerfettoDsImpl* ds_impl, void* user_arg,
+                              void* inst_ctx) -> void {
+        auto* datasource_instance = static_cast<PerfettoDataSourceInstance*>(inst_ctx);
+        delete datasource_instance;
+    };
+
+    PerfettoDsRegister(&datasource->dataSource, datasource->dataSourceName.c_str(), params);
+}
+
+jobject nativeGetPerfettoInstanceLocked(JNIEnv* env, jclass clazz, jlong dataSourcePtr,
+                                        PerfettoDsInstanceIndex instance_idx) {
+    sp<PerfettoDataSource> datasource = reinterpret_cast<PerfettoDataSource*>(dataSourcePtr);
+    auto* datasource_instance = static_cast<PerfettoDataSourceInstance*>(
+            PerfettoDsImplGetInstanceLocked(datasource->dataSource.impl, instance_idx));
+
+    if (datasource_instance == nullptr) {
+        // datasource instance doesn't exist
+        return nullptr;
+    }
+
+    return datasource_instance->GetJavaDataSourceInstance();
+}
+
+void nativeReleasePerfettoInstanceLocked(JNIEnv* env, jclass clazz, jlong dataSourcePtr,
+                                         PerfettoDsInstanceIndex instance_idx) {
+    sp<PerfettoDataSource> datasource = reinterpret_cast<PerfettoDataSource*>(dataSourcePtr);
+    PerfettoDsImplReleaseInstanceLocked(datasource->dataSource.impl, instance_idx);
+}
+
+const JNINativeMethod gMethods[] = {
+        /* name, signature, funcPtr */
+        {"nativeCreate", "(Landroid/tracing/perfetto/DataSource;Ljava/lang/String;)J",
+         (void*)nativeCreate},
+        {"nativeTrace", "(JLandroid/tracing/perfetto/TraceFunction;)V", (void*)nativeTrace},
+        {"nativeFlushAll", "(J)V", (void*)nativeFlushAll},
+        {"nativeGetFinalizer", "()J", (void*)nativeGetFinalizer},
+        {"nativeRegisterDataSource", "(JI)V", (void*)nativeRegisterDataSource},
+        {"nativeGetPerfettoInstanceLocked", "(JI)Landroid/tracing/perfetto/DataSourceInstance;",
+         (void*)nativeGetPerfettoInstanceLocked},
+        {"nativeReleasePerfettoInstanceLocked", "(JI)V",
+         (void*)nativeReleasePerfettoInstanceLocked},
+};
+
+const JNINativeMethod gMethodsTracingContext[] = {
+        /* name, signature, funcPtr */
+        {"nativeFlush", "(Landroid/tracing/perfetto/TracingContext;J)V", (void*)nativeFlush},
+};
+
+int register_android_tracing_PerfettoDataSource(JNIEnv* env) {
+    int res = jniRegisterNativeMethods(env, "android/tracing/perfetto/DataSource", gMethods,
+                                       NELEM(gMethods));
+
+    LOG_ALWAYS_FATAL_IF(res < 0, "Unable to register native methods.");
+
+    res = jniRegisterNativeMethods(env, "android/tracing/perfetto/TracingContext",
+                                   gMethodsTracingContext, NELEM(gMethodsTracingContext));
+
+    LOG_ALWAYS_FATAL_IF(res < 0, "Unable to register native methods.");
+
+    if (env->GetJavaVM(&gVm) != JNI_OK) {
+        LOG_ALWAYS_FATAL("Failed to get JavaVM from JNIEnv: %p", env);
+    }
+
+    jclass clazz = env->FindClass("android/tracing/perfetto/DataSource");
+    gPerfettoDataSourceClassInfo.clazz = MakeGlobalRefOrDie(env, clazz);
+    gPerfettoDataSourceClassInfo.createInstance =
+            env->GetMethodID(gPerfettoDataSourceClassInfo.clazz, "createInstance",
+                             "([BI)Landroid/tracing/perfetto/DataSourceInstance;");
+    gPerfettoDataSourceClassInfo.createTlsState =
+            env->GetMethodID(gPerfettoDataSourceClassInfo.clazz, "createTlsState",
+                             "(Landroid/tracing/perfetto/CreateTlsStateArgs;)Ljava/lang/Object;");
+    gPerfettoDataSourceClassInfo.createIncrementalState =
+            env->GetMethodID(gPerfettoDataSourceClassInfo.clazz, "createIncrementalState",
+                             "(Landroid/tracing/perfetto/CreateIncrementalStateArgs;)Ljava/lang/"
+                             "Object;");
+
+    clazz = env->FindClass("android/tracing/perfetto/TracingContext");
+    gTracingContextClassInfo.clazz = MakeGlobalRefOrDie(env, clazz);
+    gTracingContextClassInfo.init = env->GetMethodID(gTracingContextClassInfo.clazz, "<init>",
+                                                     "(JLjava/lang/Object;Ljava/lang/Object;)V");
+    gTracingContextClassInfo.getAndClearAllPendingTracePackets =
+            env->GetMethodID(gTracingContextClassInfo.clazz, "getAndClearAllPendingTracePackets",
+                             "()[[B");
+
+    clazz = env->FindClass("android/tracing/perfetto/CreateTlsStateArgs");
+    gCreateTlsStateArgsClassInfo.clazz = MakeGlobalRefOrDie(env, clazz);
+    gCreateTlsStateArgsClassInfo.init =
+            env->GetMethodID(gCreateTlsStateArgsClassInfo.clazz, "<init>",
+                             "(Landroid/tracing/perfetto/DataSource;I)V");
+
+    clazz = env->FindClass("android/tracing/perfetto/CreateIncrementalStateArgs");
+    gCreateIncrementalStateArgsClassInfo.clazz = MakeGlobalRefOrDie(env, clazz);
+    gCreateIncrementalStateArgsClassInfo.init =
+            env->GetMethodID(gCreateIncrementalStateArgsClassInfo.clazz, "<init>",
+                             "(Landroid/tracing/perfetto/DataSource;I)V");
+
+    return 0;
+}
+
+} // namespace android
\ No newline at end of file
diff --git a/core/jni/android_tracing_PerfettoDataSource.h b/core/jni/android_tracing_PerfettoDataSource.h
new file mode 100644
index 0000000..4ddf1d8
--- /dev/null
+++ b/core/jni/android_tracing_PerfettoDataSource.h
@@ -0,0 +1,59 @@
+/*
+ * Copyright 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.
+ */
+
+#define LOG_TAG "Perfetto"
+
+#include <android_runtime/AndroidRuntime.h>
+#include <android_runtime/Log.h>
+#include <nativehelper/JNIHelp.h>
+#include <perfetto/public/data_source.h>
+#include <perfetto/public/producer.h>
+#include <perfetto/public/protos/trace/test_event.pzc.h>
+#include <perfetto/public/protos/trace/trace_packet.pzc.h>
+#include <perfetto/tracing.h>
+#include <utils/Log.h>
+#include <utils/RefBase.h>
+
+#include <sstream>
+#include <thread>
+
+#include "android_tracing_PerfettoDataSourceInstance.h"
+#include "core_jni_helpers.h"
+
+namespace android {
+
+class PerfettoDataSource : public virtual RefBase {
+public:
+    const std::string dataSourceName;
+    struct PerfettoDs dataSource = PERFETTO_DS_INIT();
+
+    PerfettoDataSource(JNIEnv* env, jobject java_data_source, std::string data_source_name);
+    ~PerfettoDataSource();
+
+    jobject newInstance(JNIEnv* env, void* ds_config, size_t ds_config_size,
+                        PerfettoDsInstanceIndex inst_id);
+
+    jobject createTlsStateGlobalRef(JNIEnv* env, PerfettoDsInstanceIndex inst_id);
+    jobject createIncrementalStateGlobalRef(JNIEnv* env, PerfettoDsInstanceIndex inst_id);
+    void trace(JNIEnv* env, jobject trace_function);
+    void flushAll();
+
+private:
+    jobject mJavaDataSource;
+    std::map<PerfettoDsInstanceIndex, PerfettoDataSourceInstance*> mInstances;
+};
+
+} // namespace android
\ No newline at end of file
diff --git a/core/jni/android_tracing_PerfettoDataSourceInstance.cpp b/core/jni/android_tracing_PerfettoDataSourceInstance.cpp
new file mode 100644
index 0000000..e659bf1
--- /dev/null
+++ b/core/jni/android_tracing_PerfettoDataSourceInstance.cpp
@@ -0,0 +1,138 @@
+/*
+ * Copyright 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.
+ */
+
+#define LOG_TAG "Perfetto"
+
+#include "android_tracing_PerfettoDataSourceInstance.h"
+
+#include <android_runtime/AndroidRuntime.h>
+#include <android_runtime/Log.h>
+#include <nativehelper/JNIHelp.h>
+#include <perfetto/public/data_source.h>
+#include <perfetto/public/producer.h>
+#include <perfetto/public/protos/trace/test_event.pzc.h>
+#include <perfetto/public/protos/trace/trace_packet.pzc.h>
+#include <perfetto/tracing.h>
+#include <utils/Log.h>
+#include <utils/RefBase.h>
+
+#include <sstream>
+#include <thread>
+
+#include "core_jni_helpers.h"
+
+namespace android {
+
+static struct {
+    jclass clazz;
+    jmethodID init;
+} gStartCallbackArgumentsClassInfo;
+
+static struct {
+    jclass clazz;
+    jmethodID init;
+} gFlushCallbackArgumentsClassInfo;
+
+static struct {
+    jclass clazz;
+    jmethodID init;
+} gStopCallbackArgumentsClassInfo;
+
+static JavaVM* gVm;
+
+void callJavaMethodWithArgsObject(JNIEnv* env, jobject classRef, jmethodID method, jobject args) {
+    ScopedLocalRef<jobject> localClassRef(env, env->NewLocalRef(classRef));
+
+    if (localClassRef == nullptr) {
+        ALOGE("Weak reference went out of scope");
+        return;
+    }
+
+    env->CallVoidMethod(localClassRef.get(), method, args);
+
+    if (env->ExceptionCheck()) {
+        env->ExceptionDescribe();
+        LOGE_EX(env);
+        env->ExceptionClear();
+    }
+}
+
+PerfettoDataSourceInstance::PerfettoDataSourceInstance(JNIEnv* env, jobject javaDataSourceInstance,
+                                                       PerfettoDsInstanceIndex inst_idx)
+      : inst_idx(inst_idx), mJavaDataSourceInstance(env->NewGlobalRef(javaDataSourceInstance)) {}
+
+PerfettoDataSourceInstance::~PerfettoDataSourceInstance() {
+    JNIEnv* env = GetOrAttachJNIEnvironment(gVm, JNI_VERSION_1_6);
+    env->DeleteGlobalRef(mJavaDataSourceInstance);
+}
+
+void PerfettoDataSourceInstance::onStart(JNIEnv* env) {
+    ScopedLocalRef<jobject> args(env,
+                                 env->NewObject(gStartCallbackArgumentsClassInfo.clazz,
+                                                gStartCallbackArgumentsClassInfo.init));
+    jclass cls = env->GetObjectClass(mJavaDataSourceInstance);
+    jmethodID mid = env->GetMethodID(cls, "onStart",
+                                     "(Landroid/tracing/perfetto/StartCallbackArguments;)V");
+
+    callJavaMethodWithArgsObject(env, mJavaDataSourceInstance, mid, args.get());
+}
+
+void PerfettoDataSourceInstance::onFlush(JNIEnv* env) {
+    ScopedLocalRef<jobject> args(env,
+                                 env->NewObject(gFlushCallbackArgumentsClassInfo.clazz,
+                                                gFlushCallbackArgumentsClassInfo.init));
+    jclass cls = env->GetObjectClass(mJavaDataSourceInstance);
+    jmethodID mid = env->GetMethodID(cls, "onFlush",
+                                     "(Landroid/tracing/perfetto/FlushCallbackArguments;)V");
+
+    callJavaMethodWithArgsObject(env, mJavaDataSourceInstance, mid, args.get());
+}
+
+void PerfettoDataSourceInstance::onStop(JNIEnv* env) {
+    ScopedLocalRef<jobject> args(env,
+                                 env->NewObject(gStopCallbackArgumentsClassInfo.clazz,
+                                                gStopCallbackArgumentsClassInfo.init));
+    jclass cls = env->GetObjectClass(mJavaDataSourceInstance);
+    jmethodID mid =
+            env->GetMethodID(cls, "onStop", "(Landroid/tracing/perfetto/StopCallbackArguments;)V");
+
+    callJavaMethodWithArgsObject(env, mJavaDataSourceInstance, mid, args.get());
+}
+
+int register_android_tracing_PerfettoDataSourceInstance(JNIEnv* env) {
+    if (env->GetJavaVM(&gVm) != JNI_OK) {
+        LOG_ALWAYS_FATAL("Failed to get JavaVM from JNIEnv: %p", env);
+    }
+
+    jclass clazz = env->FindClass("android/tracing/perfetto/StartCallbackArguments");
+    gStartCallbackArgumentsClassInfo.clazz = MakeGlobalRefOrDie(env, clazz);
+    gStartCallbackArgumentsClassInfo.init =
+            env->GetMethodID(gStartCallbackArgumentsClassInfo.clazz, "<init>", "()V");
+
+    clazz = env->FindClass("android/tracing/perfetto/FlushCallbackArguments");
+    gFlushCallbackArgumentsClassInfo.clazz = MakeGlobalRefOrDie(env, clazz);
+    gFlushCallbackArgumentsClassInfo.init =
+            env->GetMethodID(gFlushCallbackArgumentsClassInfo.clazz, "<init>", "()V");
+
+    clazz = env->FindClass("android/tracing/perfetto/StopCallbackArguments");
+    gStopCallbackArgumentsClassInfo.clazz = MakeGlobalRefOrDie(env, clazz);
+    gStopCallbackArgumentsClassInfo.init =
+            env->GetMethodID(gStopCallbackArgumentsClassInfo.clazz, "<init>", "()V");
+
+    return 0;
+}
+
+} // namespace android
\ No newline at end of file
diff --git a/core/jni/android_tracing_PerfettoDataSourceInstance.h b/core/jni/android_tracing_PerfettoDataSourceInstance.h
new file mode 100644
index 0000000..d577655
--- /dev/null
+++ b/core/jni/android_tracing_PerfettoDataSourceInstance.h
@@ -0,0 +1,59 @@
+/*
+ * Copyright 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.
+ */
+
+#define LOG_TAG "Perfetto"
+
+#include <android_runtime/AndroidRuntime.h>
+#include <android_runtime/Log.h>
+#include <nativehelper/JNIHelp.h>
+#include <perfetto/public/data_source.h>
+#include <perfetto/public/producer.h>
+#include <perfetto/public/protos/trace/test_event.pzc.h>
+#include <perfetto/public/protos/trace/trace_packet.pzc.h>
+#include <perfetto/tracing.h>
+#include <utils/Log.h>
+#include <utils/RefBase.h>
+
+#include <sstream>
+#include <thread>
+
+#include "core_jni_helpers.h"
+
+namespace android {
+
+class PerfettoDataSourceInstance {
+public:
+    PerfettoDataSourceInstance(JNIEnv* env, jobject javaDataSourceInstance,
+                               PerfettoDsInstanceIndex inst_idx);
+    ~PerfettoDataSourceInstance();
+
+    void onStart(JNIEnv* env);
+    void onFlush(JNIEnv* env);
+    void onStop(JNIEnv* env);
+
+    jobject GetJavaDataSourceInstance() {
+        return mJavaDataSourceInstance;
+    }
+
+    PerfettoDsInstanceIndex getIndex() {
+        return inst_idx;
+    }
+
+private:
+    PerfettoDsInstanceIndex inst_idx;
+    jobject mJavaDataSourceInstance;
+};
+} // namespace android
\ No newline at end of file
diff --git a/core/jni/android_tracing_PerfettoProducer.cpp b/core/jni/android_tracing_PerfettoProducer.cpp
new file mode 100644
index 0000000..ce72f58
--- /dev/null
+++ b/core/jni/android_tracing_PerfettoProducer.cpp
@@ -0,0 +1,58 @@
+/*
+ * Copyright 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.
+ */
+
+#define LOG_TAG "Perfetto"
+
+#include <android_runtime/AndroidRuntime.h>
+#include <android_runtime/Log.h>
+#include <nativehelper/JNIHelp.h>
+#include <perfetto/public/data_source.h>
+#include <perfetto/public/producer.h>
+#include <perfetto/public/protos/trace/test_event.pzc.h>
+#include <perfetto/public/protos/trace/trace_packet.pzc.h>
+#include <perfetto/tracing.h>
+#include <utils/Log.h>
+#include <utils/RefBase.h>
+
+#include <sstream>
+#include <thread>
+
+#include "android_tracing_PerfettoDataSource.h"
+#include "core_jni_helpers.h"
+
+namespace android {
+
+void perfettoProducerInit(JNIEnv* env, jclass clazz, int backends) {
+    struct PerfettoProducerInitArgs args = PERFETTO_PRODUCER_INIT_ARGS_INIT();
+    args.backends = (PerfettoBackendTypes)backends;
+    PerfettoProducerInit(args);
+}
+
+const JNINativeMethod gMethods[] = {
+        /* name, signature, funcPtr */
+        {"nativePerfettoProducerInit", "(I)V", (void*)perfettoProducerInit},
+};
+
+int register_android_tracing_PerfettoProducer(JNIEnv* env) {
+    int res = jniRegisterNativeMethods(env, "android/tracing/perfetto/Producer", gMethods,
+                                       NELEM(gMethods));
+
+    LOG_ALWAYS_FATAL_IF(res < 0, "Unable to register native methods.");
+
+    return 0;
+}
+
+} // namespace android
\ No newline at end of file
diff --git a/core/tests/coretests/Android.bp b/core/tests/coretests/Android.bp
index c058174..531756e 100644
--- a/core/tests/coretests/Android.bp
+++ b/core/tests/coretests/Android.bp
@@ -83,6 +83,10 @@
         "com.android.text.flags-aconfig-java",
         "flag-junit",
         "ravenwood-junit",
+        "perfetto_trace_java_protos",
+        "flickerlib-parsers",
+        "flickerlib-trace_processor_shell",
+        "mockito-target-extended-minus-junit4",
     ],
 
     libs: [
diff --git a/core/tests/coretests/src/android/tracing/perfetto/DataSourceTest.java b/core/tests/coretests/src/android/tracing/perfetto/DataSourceTest.java
new file mode 100644
index 0000000..bd2f36f
--- /dev/null
+++ b/core/tests/coretests/src/android/tracing/perfetto/DataSourceTest.java
@@ -0,0 +1,664 @@
+/*
+ * 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 android.tracing.perfetto;
+
+import static android.internal.perfetto.protos.PerfettoTrace.TestEvent.PAYLOAD;
+import static android.internal.perfetto.protos.PerfettoTrace.TestEvent.TestPayload.SINGLE_INT;
+import static android.internal.perfetto.protos.PerfettoTrace.TracePacket.FOR_TESTING;
+
+import static java.io.File.createTempFile;
+import static java.nio.file.Files.createTempDirectory;
+
+import android.internal.perfetto.protos.PerfettoTrace;
+import android.tools.common.ScenarioBuilder;
+import android.tools.common.Tag;
+import android.tools.common.io.TraceType;
+import android.tools.device.traces.TraceConfig;
+import android.tools.device.traces.TraceConfigs;
+import android.tools.device.traces.io.ResultReader;
+import android.tools.device.traces.io.ResultWriter;
+import android.tools.device.traces.monitors.PerfettoTraceMonitor;
+import android.tools.device.traces.monitors.TraceMonitor;
+import android.util.proto.ProtoInputStream;
+import android.util.proto.ProtoOutputStream;
+
+import androidx.test.runner.AndroidJUnit4;
+
+import com.google.common.truth.Truth;
+import com.google.protobuf.InvalidProtocolBufferException;
+
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+
+import perfetto.protos.PerfettoConfig;
+import perfetto.protos.TracePacketOuterClass;
+
+@RunWith(AndroidJUnit4.class)
+public class DataSourceTest {
+    private final File mTracingDirectory = createTempDirectory("temp").toFile();
+
+    private final ResultWriter mWriter = new ResultWriter()
+            .forScenario(new ScenarioBuilder()
+                    .forClass(createTempFile("temp", "").getName()).build())
+            .withOutputDir(mTracingDirectory)
+            .setRunComplete();
+
+    private final TraceConfigs mTraceConfig = new TraceConfigs(
+            new TraceConfig(false, true, false),
+            new TraceConfig(false, true, false),
+            new TraceConfig(false, true, false),
+            new TraceConfig(false, true, false)
+    );
+
+    private static TestDataSource sTestDataSource;
+
+    private static TestDataSource.DataSourceInstanceProvider sInstanceProvider;
+    private static TestDataSource.TlsStateProvider sTlsStateProvider;
+    private static TestDataSource.IncrementalStateProvider sIncrementalStateProvider;
+
+    public DataSourceTest() throws IOException {}
+
+    @BeforeClass
+    public static void beforeAll() {
+        Producer.init(InitArguments.DEFAULTS);
+        setupProviders();
+        sTestDataSource = new TestDataSource(
+                (ds, idx, configStream) -> sInstanceProvider.provide(ds, idx, configStream),
+                args -> sTlsStateProvider.provide(args),
+                args -> sIncrementalStateProvider.provide(args));
+        sTestDataSource.register(DataSourceParams.DEFAULTS);
+    }
+
+    private static void setupProviders() {
+        sInstanceProvider = (ds, idx, configStream) ->
+                new TestDataSource.TestDataSourceInstance(ds, idx);
+        sTlsStateProvider = args -> new TestDataSource.TestTlsState();
+        sIncrementalStateProvider = args -> new TestDataSource.TestIncrementalState();
+    }
+
+    @Before
+    public void setup() {
+        setupProviders();
+    }
+
+    @Test
+    public void canTraceData() throws InvalidProtocolBufferException {
+        final TraceMonitor traceMonitor = PerfettoTraceMonitor.newBuilder()
+                .enableCustomTrace(PerfettoConfig.DataSourceConfig.newBuilder()
+                        .setName(sTestDataSource.name).build()).build();
+
+        try {
+            traceMonitor.start();
+
+            sTestDataSource.trace((ctx) -> {
+                final ProtoOutputStream protoOutputStream = ctx.newTracePacket();
+                long forTestingToken = protoOutputStream.start(FOR_TESTING);
+                long payloadToken = protoOutputStream.start(PAYLOAD);
+                protoOutputStream.write(SINGLE_INT, 10);
+                protoOutputStream.end(payloadToken);
+                protoOutputStream.end(forTestingToken);
+            });
+        } finally {
+            traceMonitor.stop(mWriter);
+        }
+
+        final ResultReader reader = new ResultReader(mWriter.write(), mTraceConfig);
+        final byte[] rawProtoFromFile = reader.readBytes(TraceType.PERFETTO, Tag.ALL);
+        assert rawProtoFromFile != null;
+        final perfetto.protos.TraceOuterClass.Trace trace = perfetto.protos.TraceOuterClass.Trace
+                .parseFrom(rawProtoFromFile);
+
+        Truth.assertThat(trace.getPacketCount()).isGreaterThan(0);
+        final List<TracePacketOuterClass.TracePacket> tracePackets = trace.getPacketList()
+                .stream().filter(TracePacketOuterClass.TracePacket::hasForTesting).toList();
+        final List<TracePacketOuterClass.TracePacket> matchingPackets = tracePackets.stream()
+                .filter(it -> it.getForTesting().getPayload().getSingleInt() == 10).toList();
+        Truth.assertThat(matchingPackets).hasSize(1);
+    }
+
+    @Test
+    public void canUseTlsStateForCustomState() {
+        final int expectedStateTestValue = 10;
+        final AtomicInteger actualStateTestValue = new AtomicInteger();
+
+        final TraceMonitor traceMonitor = PerfettoTraceMonitor.newBuilder()
+                .enableCustomTrace(PerfettoConfig.DataSourceConfig.newBuilder()
+                        .setName(sTestDataSource.name).build()).build();
+
+        try {
+            traceMonitor.start();
+
+            sTestDataSource.trace((ctx) -> {
+                TestDataSource.TestTlsState state = ctx.getCustomTlsState();
+                state.testStateValue = expectedStateTestValue;
+            });
+
+            sTestDataSource.trace((ctx) -> {
+                TestDataSource.TestTlsState state = ctx.getCustomTlsState();
+                actualStateTestValue.set(state.testStateValue);
+            });
+        } finally {
+            traceMonitor.stop(mWriter);
+        }
+
+        Truth.assertThat(actualStateTestValue.get()).isEqualTo(expectedStateTestValue);
+    }
+
+    @Test
+    public void eachInstanceHasOwnTlsState() {
+        final int[] expectedStateTestValues = new int[] { 1, 2 };
+        final int[] actualStateTestValues = new int[] { 0, 0 };
+
+        final TraceMonitor traceMonitor1 = PerfettoTraceMonitor.newBuilder()
+                .enableCustomTrace(PerfettoConfig.DataSourceConfig.newBuilder()
+                        .setName(sTestDataSource.name).build()).build();
+        final TraceMonitor traceMonitor2 = PerfettoTraceMonitor.newBuilder()
+                .enableCustomTrace(PerfettoConfig.DataSourceConfig.newBuilder()
+                        .setName(sTestDataSource.name).build()).build();
+
+        try {
+            traceMonitor1.start();
+            try {
+                traceMonitor2.start();
+
+                AtomicInteger index = new AtomicInteger(0);
+                sTestDataSource.trace((ctx) -> {
+                    TestDataSource.TestTlsState state = ctx.getCustomTlsState();
+                    state.testStateValue = expectedStateTestValues[index.getAndIncrement()];
+                });
+
+                index.set(0);
+                sTestDataSource.trace((ctx) -> {
+                    TestDataSource.TestTlsState state = ctx.getCustomTlsState();
+                    actualStateTestValues[index.getAndIncrement()] = state.testStateValue;
+                });
+            } finally {
+                traceMonitor1.stop(mWriter);
+            }
+        } finally {
+            traceMonitor2.stop(mWriter);
+        }
+
+        Truth.assertThat(actualStateTestValues[0]).isEqualTo(expectedStateTestValues[0]);
+        Truth.assertThat(actualStateTestValues[1]).isEqualTo(expectedStateTestValues[1]);
+    }
+
+    @Test
+    public void eachThreadHasOwnTlsState() throws InterruptedException {
+        final int thread1ExpectedStateValue = 1;
+        final int thread2ExpectedStateValue = 2;
+
+        final AtomicInteger thread1ActualStateValue = new AtomicInteger();
+        final AtomicInteger thread2ActualStateValue = new AtomicInteger();
+
+        final CountDownLatch setUpLatch = new CountDownLatch(2);
+        final CountDownLatch setStateLatch = new CountDownLatch(2);
+        final CountDownLatch setOutStateLatch = new CountDownLatch(2);
+
+        final RunnableCreator createTask = (stateValue, stateOut) -> () -> {
+            Producer.init(InitArguments.DEFAULTS);
+
+            setUpLatch.countDown();
+
+            try {
+                setUpLatch.await(3, TimeUnit.SECONDS);
+            } catch (InterruptedException e) {
+                throw new RuntimeException(e);
+            }
+
+            sTestDataSource.trace((ctx) -> {
+                TestDataSource.TestTlsState state = ctx.getCustomTlsState();
+                state.testStateValue = stateValue;
+                setStateLatch.countDown();
+            });
+
+            try {
+                setStateLatch.await(3, TimeUnit.SECONDS);
+            } catch (InterruptedException e) {
+                throw new RuntimeException(e);
+            }
+
+            sTestDataSource.trace((ctx) -> {
+                stateOut.set(ctx.getCustomTlsState().testStateValue);
+                setOutStateLatch.countDown();
+            });
+        };
+
+        final TraceMonitor traceMonitor = PerfettoTraceMonitor.newBuilder()
+                .enableCustomTrace(PerfettoConfig.DataSourceConfig.newBuilder()
+                        .setName(sTestDataSource.name).build()).build();
+
+        try {
+            traceMonitor.start();
+
+            new Thread(
+                    createTask.create(thread1ExpectedStateValue, thread1ActualStateValue)).start();
+            new Thread(
+                    createTask.create(thread2ExpectedStateValue, thread2ActualStateValue)).start();
+
+            setOutStateLatch.await(3, TimeUnit.SECONDS);
+
+        } finally {
+            traceMonitor.stop(mWriter);
+        }
+
+        Truth.assertThat(thread1ActualStateValue.get()).isEqualTo(thread1ExpectedStateValue);
+        Truth.assertThat(thread2ActualStateValue.get()).isEqualTo(thread2ExpectedStateValue);
+    }
+
+    @Test
+    public void incrementalStateIsReset() throws InterruptedException {
+
+        final TraceMonitor traceMonitor = PerfettoTraceMonitor.newBuilder()
+                .enableCustomTrace(PerfettoConfig.DataSourceConfig.newBuilder()
+                        .setName(sTestDataSource.name).build())
+                .setIncrementalTimeout(10)
+                .build();
+
+        final AtomicInteger testStateValue = new AtomicInteger();
+        try {
+            traceMonitor.start();
+
+            sTestDataSource.trace(ctx -> ctx.getIncrementalState().testStateValue = 1);
+
+            // Timeout to make sure the incremental state is cleared.
+            Thread.sleep(1000);
+
+            sTestDataSource.trace(ctx ->
+                    testStateValue.set(ctx.getIncrementalState().testStateValue));
+        } finally {
+            traceMonitor.stop(mWriter);
+        }
+
+        Truth.assertThat(testStateValue.get()).isNotEqualTo(1);
+    }
+
+    @Test
+    public void getInstanceConfigOnCreateInstance() throws IOException {
+        final int expectedDummyIntValue = 10;
+        AtomicReference<ProtoInputStream> configStream = new AtomicReference<>();
+        sInstanceProvider = (ds, idx, config) -> {
+            configStream.set(config);
+            return new TestDataSource.TestDataSourceInstance(ds, idx);
+        };
+
+        final TraceMonitor monitor = PerfettoTraceMonitor.newBuilder()
+                .enableCustomTrace(PerfettoConfig.DataSourceConfig.newBuilder()
+                        .setName(sTestDataSource.name)
+                        .setForTesting(PerfettoConfig.TestConfig.newBuilder().setDummyFields(
+                                PerfettoConfig.TestConfig.DummyFields.newBuilder()
+                                        .setFieldInt32(expectedDummyIntValue)
+                                        .build())
+                                .build())
+                        .build())
+                .build();
+
+        try {
+            monitor.start();
+        } finally {
+            monitor.stop(mWriter);
+        }
+
+        int configDummyIntValue = 0;
+        while (configStream.get().nextField() != ProtoInputStream.NO_MORE_FIELDS) {
+            if (configStream.get().getFieldNumber()
+                    == (int) PerfettoTrace.DataSourceConfig.FOR_TESTING) {
+                final long forTestingToken = configStream.get()
+                        .start(PerfettoTrace.DataSourceConfig.FOR_TESTING);
+                while (configStream.get().nextField() != ProtoInputStream.NO_MORE_FIELDS) {
+                    if (configStream.get().getFieldNumber()
+                            == (int) PerfettoTrace.TestConfig.DUMMY_FIELDS) {
+                        final long dummyFieldsToken = configStream.get()
+                                .start(PerfettoTrace.TestConfig.DUMMY_FIELDS);
+                        while (configStream.get().nextField() != ProtoInputStream.NO_MORE_FIELDS) {
+                            if (configStream.get().getFieldNumber()
+                                    == (int) PerfettoTrace.TestConfig.DummyFields.FIELD_INT32) {
+                                int val = configStream.get().readInt(
+                                        PerfettoTrace.TestConfig.DummyFields.FIELD_INT32);
+                                if (val != 0) {
+                                    configDummyIntValue = val;
+                                    break;
+                                }
+                            }
+                        }
+                        configStream.get().end(dummyFieldsToken);
+                        break;
+                    }
+                }
+                configStream.get().end(forTestingToken);
+                break;
+            }
+        }
+
+        Truth.assertThat(configDummyIntValue).isEqualTo(expectedDummyIntValue);
+    }
+
+    @Test
+    public void multipleTraceInstances() throws IOException, InterruptedException {
+        final int instanceCount = 3;
+
+        final List<TraceMonitor> monitors = new ArrayList<>();
+        final List<ResultWriter> writers = new ArrayList<>();
+
+        for (int i = 0; i < instanceCount; i++) {
+            final ResultWriter writer = new ResultWriter()
+                    .forScenario(new ScenarioBuilder()
+                            .forClass(createTempFile("temp", "").getName()).build())
+                    .withOutputDir(mTracingDirectory)
+                    .setRunComplete();
+            writers.add(writer);
+        }
+
+        // Start at 1 because 0 is considered null value so payload will be ignored in that case
+        TestDataSource.TestTlsState.lastIndex = 1;
+
+        final AtomicInteger traceCallCount = new AtomicInteger();
+        final CountDownLatch latch = new CountDownLatch(instanceCount);
+
+        try {
+            // Start instances
+            for (int i = 0; i < instanceCount; i++) {
+                final TraceMonitor traceMonitor = PerfettoTraceMonitor.newBuilder()
+                        .enableCustomTrace(PerfettoConfig.DataSourceConfig.newBuilder()
+                                .setName(sTestDataSource.name).build()).build();
+                monitors.add(traceMonitor);
+                traceMonitor.start();
+            }
+
+            // Trace the stateIndex of the tracing instance.
+            sTestDataSource.trace(ctx -> {
+                final int testIntValue = ctx.getCustomTlsState().stateIndex;
+                traceCallCount.incrementAndGet();
+
+                final ProtoOutputStream os = ctx.newTracePacket();
+                long forTestingToken = os.start(FOR_TESTING);
+                long payloadToken = os.start(PAYLOAD);
+                os.write(SINGLE_INT, testIntValue);
+                os.end(payloadToken);
+                os.end(forTestingToken);
+
+                latch.countDown();
+            });
+        } finally {
+            // Stop instances
+            for (int i = 0; i < instanceCount; i++) {
+                final TraceMonitor monitor = monitors.get(i);
+                final ResultWriter writer = writers.get(i);
+                monitor.stop(writer);
+            }
+        }
+
+        latch.await(3, TimeUnit.SECONDS);
+        Truth.assertThat(traceCallCount.get()).isEqualTo(instanceCount);
+
+        for (int i = 0; i < instanceCount; i++) {
+            final int expectedTracedValue = i + 1;
+            final ResultWriter writer = writers.get(i);
+            final ResultReader reader = new ResultReader(writer.write(), mTraceConfig);
+            final byte[] rawProtoFromFile = reader.readBytes(TraceType.PERFETTO, Tag.ALL);
+            assert rawProtoFromFile != null;
+            final perfetto.protos.TraceOuterClass.Trace trace =
+                    perfetto.protos.TraceOuterClass.Trace.parseFrom(rawProtoFromFile);
+
+            Truth.assertThat(trace.getPacketCount()).isGreaterThan(0);
+            final List<TracePacketOuterClass.TracePacket> tracePackets = trace.getPacketList()
+                    .stream().filter(TracePacketOuterClass.TracePacket::hasForTesting).toList();
+            Truth.assertWithMessage("One packet has for testing data")
+                    .that(tracePackets).hasSize(1);
+
+            final List<TracePacketOuterClass.TracePacket> matchingPackets =
+                    tracePackets.stream()
+                            .filter(it -> it.getForTesting().getPayload()
+                                    .getSingleInt() == expectedTracedValue).toList();
+            Truth.assertWithMessage(
+                            "One packet has testing data with a payload with the expected value")
+                    .that(matchingPackets).hasSize(1);
+        }
+    }
+
+    @Test
+    public void onStartCallbackTriggered() throws InterruptedException {
+        final CountDownLatch latch = new CountDownLatch(1);
+
+        final AtomicBoolean callbackCalled = new AtomicBoolean(false);
+        sInstanceProvider = (ds, idx, config) -> new TestDataSource.TestDataSourceInstance(
+                        ds,
+                        idx,
+                        (args) -> {
+                            callbackCalled.set(true);
+                            latch.countDown();
+                        },
+                        (args) -> {},
+                        (args) -> {}
+        );
+
+        final TraceMonitor traceMonitor = PerfettoTraceMonitor.newBuilder()
+                .enableCustomTrace(PerfettoConfig.DataSourceConfig.newBuilder()
+                        .setName(sTestDataSource.name).build()).build();
+
+        Truth.assertThat(callbackCalled.get()).isFalse();
+        try {
+            traceMonitor.start();
+            latch.await(3, TimeUnit.SECONDS);
+            Truth.assertThat(callbackCalled.get()).isTrue();
+        } finally {
+            traceMonitor.stop(mWriter);
+        }
+    }
+
+    @Test
+    public void onFlushCallbackTriggered() throws InterruptedException {
+        final CountDownLatch latch = new CountDownLatch(1);
+        final AtomicBoolean callbackCalled = new AtomicBoolean(false);
+        sInstanceProvider = (ds, idx, config) ->
+                new TestDataSource.TestDataSourceInstance(
+                        ds,
+                        idx,
+                        (args) -> {},
+                        (args) -> {
+                            callbackCalled.set(true);
+                            latch.countDown();
+                        },
+                        (args) -> {}
+                );
+
+        final TraceMonitor traceMonitor = PerfettoTraceMonitor.newBuilder()
+                .enableCustomTrace(PerfettoConfig.DataSourceConfig.newBuilder()
+                        .setName(sTestDataSource.name).build()).build();
+
+        try {
+            traceMonitor.start();
+            Truth.assertThat(callbackCalled.get()).isFalse();
+            sTestDataSource.trace((ctx) -> {
+                final ProtoOutputStream protoOutputStream = ctx.newTracePacket();
+                long forTestingToken = protoOutputStream.start(FOR_TESTING);
+                long payloadToken = protoOutputStream.start(PAYLOAD);
+                protoOutputStream.write(SINGLE_INT, 10);
+                protoOutputStream.end(payloadToken);
+                protoOutputStream.end(forTestingToken);
+            });
+        } finally {
+            traceMonitor.stop(mWriter);
+        }
+
+        latch.await(3, TimeUnit.SECONDS);
+        Truth.assertThat(callbackCalled.get()).isTrue();
+    }
+
+    @Test
+    public void onStopCallbackTriggered() throws InterruptedException {
+        final CountDownLatch latch = new CountDownLatch(1);
+        final AtomicBoolean callbackCalled = new AtomicBoolean(false);
+        sInstanceProvider = (ds, idx, config) ->
+                new TestDataSource.TestDataSourceInstance(
+                        ds,
+                        idx,
+                        (args) -> {},
+                        (args) -> {},
+                        (args) -> {
+                            callbackCalled.set(true);
+                            latch.countDown();
+                        }
+                );
+
+        final TraceMonitor traceMonitor = PerfettoTraceMonitor.newBuilder()
+                .enableCustomTrace(PerfettoConfig.DataSourceConfig.newBuilder()
+                        .setName(sTestDataSource.name).build()).build();
+
+        try {
+            traceMonitor.start();
+            Truth.assertThat(callbackCalled.get()).isFalse();
+        } finally {
+            traceMonitor.stop(mWriter);
+        }
+
+        latch.await(3, TimeUnit.SECONDS);
+        Truth.assertThat(callbackCalled.get()).isTrue();
+    }
+
+    @Test
+    public void canUseDataSourceInstanceToCreateTlsState() throws InvalidProtocolBufferException {
+        final Object testObject = new Object();
+
+        sInstanceProvider = (ds, idx, configStream) -> {
+            final TestDataSource.TestDataSourceInstance dsInstance =
+                    new TestDataSource.TestDataSourceInstance(ds, idx);
+            dsInstance.testObject = testObject;
+            return dsInstance;
+        };
+
+        sTlsStateProvider = args -> {
+            final TestDataSource.TestTlsState tlsState = new TestDataSource.TestTlsState();
+
+            try (TestDataSource.TestDataSourceInstance dataSourceInstance =
+                         args.getDataSourceInstanceLocked()) {
+                if (dataSourceInstance != null) {
+                    tlsState.testStateValue = dataSourceInstance.testObject.hashCode();
+                }
+            }
+
+            return tlsState;
+        };
+
+        final TraceMonitor traceMonitor = PerfettoTraceMonitor.newBuilder()
+                .enableCustomTrace(PerfettoConfig.DataSourceConfig.newBuilder()
+                        .setName(sTestDataSource.name).build()).build();
+
+        try {
+            traceMonitor.start();
+            sTestDataSource.trace((ctx) -> {
+                final ProtoOutputStream protoOutputStream = ctx.newTracePacket();
+                long forTestingToken = protoOutputStream.start(FOR_TESTING);
+                long payloadToken = protoOutputStream.start(PAYLOAD);
+                protoOutputStream.write(SINGLE_INT, ctx.getCustomTlsState().testStateValue);
+                protoOutputStream.end(payloadToken);
+                protoOutputStream.end(forTestingToken);
+            });
+        } finally {
+            traceMonitor.stop(mWriter);
+        }
+
+        final ResultReader reader = new ResultReader(mWriter.write(), mTraceConfig);
+        final byte[] rawProtoFromFile = reader.readBytes(TraceType.PERFETTO, Tag.ALL);
+        assert rawProtoFromFile != null;
+        final perfetto.protos.TraceOuterClass.Trace trace = perfetto.protos.TraceOuterClass.Trace
+                .parseFrom(rawProtoFromFile);
+
+        Truth.assertThat(trace.getPacketCount()).isGreaterThan(0);
+        final List<TracePacketOuterClass.TracePacket> tracePackets = trace.getPacketList()
+                .stream().filter(TracePacketOuterClass.TracePacket::hasForTesting).toList();
+        final List<TracePacketOuterClass.TracePacket> matchingPackets = tracePackets.stream()
+                .filter(it -> it.getForTesting().getPayload().getSingleInt()
+                        == testObject.hashCode()).toList();
+        Truth.assertThat(matchingPackets).hasSize(1);
+    }
+
+    @Test
+    public void canUseDataSourceInstanceToCreateIncrementalState()
+            throws InvalidProtocolBufferException {
+        final Object testObject = new Object();
+
+        sInstanceProvider = (ds, idx, configStream) -> {
+            final TestDataSource.TestDataSourceInstance dsInstance =
+                    new TestDataSource.TestDataSourceInstance(ds, idx);
+            dsInstance.testObject = testObject;
+            return dsInstance;
+        };
+
+        sIncrementalStateProvider = args -> {
+            final TestDataSource.TestIncrementalState incrementalState =
+                    new TestDataSource.TestIncrementalState();
+
+            try (TestDataSource.TestDataSourceInstance dataSourceInstance =
+                    args.getDataSourceInstanceLocked()) {
+                if (dataSourceInstance != null) {
+                    incrementalState.testStateValue = dataSourceInstance.testObject.hashCode();
+                }
+            }
+
+            return incrementalState;
+        };
+
+        final TraceMonitor traceMonitor = PerfettoTraceMonitor.newBuilder()
+                .enableCustomTrace(PerfettoConfig.DataSourceConfig.newBuilder()
+                        .setName(sTestDataSource.name).build()).build();
+
+        try {
+            traceMonitor.start();
+            sTestDataSource.trace((ctx) -> {
+                final ProtoOutputStream protoOutputStream = ctx.newTracePacket();
+                long forTestingToken = protoOutputStream.start(FOR_TESTING);
+                long payloadToken = protoOutputStream.start(PAYLOAD);
+                protoOutputStream.write(SINGLE_INT, ctx.getIncrementalState().testStateValue);
+                protoOutputStream.end(payloadToken);
+                protoOutputStream.end(forTestingToken);
+            });
+        } finally {
+            traceMonitor.stop(mWriter);
+        }
+
+        final ResultReader reader = new ResultReader(mWriter.write(), mTraceConfig);
+        final byte[] rawProtoFromFile = reader.readBytes(TraceType.PERFETTO, Tag.ALL);
+        assert rawProtoFromFile != null;
+        final perfetto.protos.TraceOuterClass.Trace trace = perfetto.protos.TraceOuterClass.Trace
+                .parseFrom(rawProtoFromFile);
+
+        Truth.assertThat(trace.getPacketCount()).isGreaterThan(0);
+        final List<TracePacketOuterClass.TracePacket> tracePackets = trace.getPacketList()
+                .stream().filter(TracePacketOuterClass.TracePacket::hasForTesting).toList();
+        final List<TracePacketOuterClass.TracePacket> matchingPackets = tracePackets.stream()
+                .filter(it -> it.getForTesting().getPayload().getSingleInt()
+                        == testObject.hashCode()).toList();
+        Truth.assertThat(matchingPackets).hasSize(1);
+    }
+
+    interface RunnableCreator {
+        Runnable create(int state, AtomicInteger stateOut);
+    }
+}
diff --git a/core/tests/coretests/src/android/tracing/perfetto/TestDataSource.java b/core/tests/coretests/src/android/tracing/perfetto/TestDataSource.java
new file mode 100644
index 0000000..d78f78b
--- /dev/null
+++ b/core/tests/coretests/src/android/tracing/perfetto/TestDataSource.java
@@ -0,0 +1,122 @@
+/*
+ * 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 android.tracing.perfetto;
+
+import android.util.proto.ProtoInputStream;
+
+import java.util.UUID;
+import java.util.function.Consumer;
+
+public class TestDataSource extends DataSource<TestDataSource.TestDataSourceInstance,
+        TestDataSource.TestTlsState, TestDataSource.TestIncrementalState> {
+    private final DataSourceInstanceProvider mDataSourceInstanceProvider;
+    private final TlsStateProvider mTlsStateProvider;
+    private final IncrementalStateProvider mIncrementalStateProvider;
+
+    interface DataSourceInstanceProvider {
+        TestDataSourceInstance provide(
+                TestDataSource dataSource, int instanceIndex, ProtoInputStream configStream);
+    }
+
+    interface TlsStateProvider {
+        TestTlsState provide(CreateTlsStateArgs<TestDataSourceInstance> args);
+    }
+
+    interface IncrementalStateProvider {
+        TestIncrementalState provide(CreateIncrementalStateArgs<TestDataSourceInstance> args);
+    }
+
+    public TestDataSource() {
+        this((ds, idx, config) -> new TestDataSourceInstance(ds, idx),
+                args -> new TestTlsState(), args -> new TestIncrementalState());
+    }
+
+    public TestDataSource(
+            DataSourceInstanceProvider dataSourceInstanceProvider,
+            TlsStateProvider tlsStateProvider,
+            IncrementalStateProvider incrementalStateProvider
+    ) {
+        super("android.tracing.perfetto.TestDataSource#" + UUID.randomUUID().toString());
+        this.mDataSourceInstanceProvider = dataSourceInstanceProvider;
+        this.mTlsStateProvider = tlsStateProvider;
+        this.mIncrementalStateProvider = incrementalStateProvider;
+    }
+
+    @Override
+    public TestDataSourceInstance createInstance(ProtoInputStream configStream, int instanceIndex) {
+        return mDataSourceInstanceProvider.provide(this, instanceIndex, configStream);
+    }
+
+    @Override
+    public TestTlsState createTlsState(CreateTlsStateArgs args) {
+        return mTlsStateProvider.provide(args);
+    }
+
+    @Override
+    public TestIncrementalState createIncrementalState(CreateIncrementalStateArgs args) {
+        return mIncrementalStateProvider.provide(args);
+    }
+
+    public static class TestTlsState {
+        public int testStateValue;
+        public int stateIndex = lastIndex++;
+
+        public static int lastIndex = 0;
+    }
+
+    public static class TestIncrementalState {
+        public int testStateValue;
+    }
+
+    public static class TestDataSourceInstance extends DataSourceInstance {
+        public Object testObject;
+        Consumer<StartCallbackArguments> mStartCallback;
+        Consumer<FlushCallbackArguments> mFlushCallback;
+        Consumer<StopCallbackArguments> mStopCallback;
+
+        public TestDataSourceInstance(DataSource dataSource, int instanceIndex) {
+            this(dataSource, instanceIndex, args -> {}, args -> {}, args -> {});
+        }
+
+        public TestDataSourceInstance(
+                DataSource dataSource,
+                int instanceIndex,
+                Consumer<StartCallbackArguments> startCallback,
+                Consumer<FlushCallbackArguments> flushCallback,
+                Consumer<StopCallbackArguments> stopCallback) {
+            super(dataSource, instanceIndex);
+            this.mStartCallback = startCallback;
+            this.mFlushCallback = flushCallback;
+            this.mStopCallback = stopCallback;
+        }
+
+        @Override
+        public void onStart(StartCallbackArguments args) {
+            this.mStartCallback.accept(args);
+        }
+
+        @Override
+        public void onFlush(FlushCallbackArguments args) {
+            this.mFlushCallback.accept(args);
+        }
+
+        @Override
+        public void onStop(StopCallbackArguments args) {
+            this.mStopCallback.accept(args);
+        }
+    }
+}
diff --git a/framework-jarjar-rules.txt b/framework-jarjar-rules.txt
index 03b268d..6339a87 100644
--- a/framework-jarjar-rules.txt
+++ b/framework-jarjar-rules.txt
@@ -8,3 +8,6 @@
 
 # for modules-utils-build dependency
 rule com.android.modules.utils.build.** android.internal.modules.utils.build.@1
+
+# For Perfetto proto dependencies
+rule perfetto.protos.** android.internal.perfetto.protos.@1
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
index 066f38b..83d555c 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
@@ -822,11 +822,6 @@
         // Checks if container should be updated before apply new parentInfo.
         final boolean shouldUpdateContainer = taskContainer.shouldUpdateContainer(parentInfo);
         taskContainer.updateTaskFragmentParentInfo(parentInfo);
-        if (!taskContainer.isVisible()) {
-            // Don't update containers if the task is not visible. We only update containers when
-            // parentInfo#isVisibleRequested is true.
-            return;
-        }
 
         // If the last direct activity of the host task is dismissed and the overlay container is
         // the only taskFragment, the overlay container should also be dismissed.
diff --git a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/OverlayPresentationTest.java b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/OverlayPresentationTest.java
index bc92101..4e7b760 100644
--- a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/OverlayPresentationTest.java
+++ b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/OverlayPresentationTest.java
@@ -475,8 +475,10 @@
     @Test
     public void testOnTaskFragmentParentInfoChanged_positionOnlyChange_earlyReturn() {
         final TaskFragmentContainer overlayContainer = createTestOverlayContainer(TASK_ID, "test");
-
         final TaskContainer taskContainer = overlayContainer.getTaskContainer();
+
+        assertThat(taskContainer.getOverlayContainer()).isEqualTo(overlayContainer);
+
         spyOn(taskContainer);
         final TaskContainer.TaskProperties taskProperties = taskContainer.getTaskProperties();
         final TaskFragmentParentInfo parentInfo = new TaskFragmentParentInfo(
@@ -495,6 +497,30 @@
                 .that(taskContainer.getOverlayContainer()).isNull();
     }
 
+    @Test
+    public void testOnTaskFragmentParentInfoChanged_invisibleTask_callDismissOverlayContainer() {
+        final TaskFragmentContainer overlayContainer = createTestOverlayContainer(TASK_ID, "test");
+        final TaskContainer taskContainer = overlayContainer.getTaskContainer();
+
+        assertThat(taskContainer.getOverlayContainer()).isEqualTo(overlayContainer);
+
+        spyOn(taskContainer);
+        final TaskContainer.TaskProperties taskProperties = taskContainer.getTaskProperties();
+        final TaskFragmentParentInfo parentInfo = new TaskFragmentParentInfo(
+                new Configuration(taskProperties.getConfiguration()), taskProperties.getDisplayId(),
+                false /* visible */, false /* hasDirectActivity */, null /* decorSurface */);
+
+        mSplitController.onTaskFragmentParentInfoChanged(mTransaction, TASK_ID, parentInfo);
+
+        // The parent info must be applied to the task container
+        verify(taskContainer).updateTaskFragmentParentInfo(parentInfo);
+        verify(mSplitController, never()).updateContainer(any(), any());
+
+        assertWithMessage("The overlay container must still be dismissed even if "
+                + "#updateContainer is not called")
+                .that(taskContainer.getOverlayContainer()).isNull();
+    }
+
     /**
      * A simplified version of {@link SplitController.ActivityStartMonitor
      * #createOrUpdateOverlayTaskFragmentIfNeeded}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/TaskSnapshotWindow.java b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/TaskSnapshotWindow.java
index e6418f3..1a0c011 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/TaskSnapshotWindow.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/TaskSnapshotWindow.java
@@ -135,7 +135,6 @@
         } catch (RemoteException e) {
             snapshotSurface.clearWindowSynced();
         }
-        window.setOuter(snapshotSurface);
         try {
             Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "TaskSnapshot#relayout");
             session.relayout(window, layoutParams, -1, -1, View.VISIBLE, 0, 0, 0,
@@ -161,7 +160,7 @@
             ShellExecutor splashScreenExecutor) {
         mSplashScreenExecutor = splashScreenExecutor;
         mSession = WindowManagerGlobal.getWindowSession();
-        mWindow = new Window();
+        mWindow = new Window(this);
         mWindow.setSession(mSession);
         int backgroundColor = taskDescription.getBackgroundColor();
         mBackgroundPaint.setColor(backgroundColor != 0 ? backgroundColor : WHITE);
@@ -204,9 +203,9 @@
     }
 
     static class Window extends BaseIWindow {
-        private WeakReference<TaskSnapshotWindow> mOuter;
+        private final WeakReference<TaskSnapshotWindow> mOuter;
 
-        public void setOuter(TaskSnapshotWindow outer) {
+        Window(TaskSnapshotWindow outer) {
             mOuter = new WeakReference<>(outer);
         }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedHandler.java
index bf783e6..8c2203e 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedHandler.java
@@ -21,17 +21,9 @@
 import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
 import static android.view.Display.DEFAULT_DISPLAY;
 import static android.view.WindowManager.TRANSIT_CHANGE;
-import static android.view.WindowManager.TRANSIT_FLAG_KEYGUARD_UNOCCLUDING;
 import static android.view.WindowManager.TRANSIT_PIP;
-import static android.view.WindowManager.TRANSIT_TO_BACK;
 import static android.window.TransitionInfo.FLAG_IN_TASK_WITH_EMBEDDED_ACTIVITY;
-import static android.window.TransitionInfo.FLAG_IS_WALLPAPER;
 
-import static com.android.wm.shell.common.split.SplitScreenConstants.FLAG_IS_DIVIDER_BAR;
-import static com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_UNDEFINED;
-import static com.android.wm.shell.pip.PipAnimationController.ANIM_TYPE_ALPHA;
-import static com.android.wm.shell.splitscreen.SplitScreen.STAGE_TYPE_UNDEFINED;
-import static com.android.wm.shell.splitscreen.SplitScreenController.EXIT_REASON_CHILD_TASK_ENTER_PIP;
 import static com.android.wm.shell.util.TransitionUtil.isOpeningType;
 
 import android.annotation.NonNull;
@@ -56,7 +48,6 @@
 import com.android.wm.shell.pip.PipTransitionController;
 import com.android.wm.shell.protolog.ShellProtoLogGroup;
 import com.android.wm.shell.recents.RecentsTransitionHandler;
-import com.android.wm.shell.splitscreen.SplitScreen;
 import com.android.wm.shell.splitscreen.SplitScreenController;
 import com.android.wm.shell.splitscreen.StageCoordinator;
 import com.android.wm.shell.sysui.ShellInit;
@@ -84,7 +75,7 @@
     private UnfoldTransitionHandler mUnfoldHandler;
     private ActivityEmbeddingController mActivityEmbeddingController;
 
-    private static class MixedTransition {
+    abstract static class MixedTransition {
         static final int TYPE_ENTER_PIP_FROM_SPLIT = 1;
 
         /** Both the display and split-state (enter/exit) is changing */
@@ -124,15 +115,11 @@
         int mAnimType = ANIM_TYPE_DEFAULT;
         final IBinder mTransition;
 
-        private final Transitions mPlayer;
-        private final DefaultMixedHandler mMixedHandler;
-        private final PipTransitionController mPipHandler;
-        private final RecentsTransitionHandler mRecentsHandler;
-        private final StageCoordinator mSplitHandler;
-        private final KeyguardTransitionHandler mKeyguardHandler;
-        private final DesktopTasksController mDesktopTasksController;
-        private final UnfoldTransitionHandler mUnfoldHandler;
-        private final ActivityEmbeddingController mActivityEmbeddingController;
+        protected final Transitions mPlayer;
+        protected final DefaultMixedHandler mMixedHandler;
+        protected final PipTransitionController mPipHandler;
+        protected final StageCoordinator mSplitHandler;
+        protected final KeyguardTransitionHandler mKeyguardHandler;
 
         Transitions.TransitionHandler mLeftoversHandler = null;
         TransitionInfo mInfo = null;
@@ -156,409 +143,33 @@
 
         MixedTransition(int type, IBinder transition, Transitions player,
                 DefaultMixedHandler mixedHandler, PipTransitionController pipHandler,
-                RecentsTransitionHandler recentsHandler, StageCoordinator splitHandler,
-                KeyguardTransitionHandler keyguardHandler,
-                DesktopTasksController desktopTasksController,
-                UnfoldTransitionHandler unfoldHandler,
-                ActivityEmbeddingController activityEmbeddingController) {
+                StageCoordinator splitHandler, KeyguardTransitionHandler keyguardHandler) {
             mType = type;
             mTransition = transition;
             mPlayer = player;
             mMixedHandler = mixedHandler;
             mPipHandler = pipHandler;
-            mRecentsHandler = recentsHandler;
             mSplitHandler = splitHandler;
             mKeyguardHandler = keyguardHandler;
-            mDesktopTasksController = desktopTasksController;
-            mUnfoldHandler = unfoldHandler;
-            mActivityEmbeddingController = activityEmbeddingController;
-
-            switch (type) {
-                case TYPE_RECENTS_DURING_DESKTOP:
-                case TYPE_RECENTS_DURING_KEYGUARD:
-                case TYPE_RECENTS_DURING_SPLIT:
-                    mLeftoversHandler = mRecentsHandler;
-                    break;
-                case TYPE_UNFOLD:
-                    mLeftoversHandler = mUnfoldHandler;
-                    break;
-                case TYPE_DISPLAY_AND_SPLIT_CHANGE:
-                case TYPE_ENTER_PIP_FROM_ACTIVITY_EMBEDDING:
-                case TYPE_ENTER_PIP_FROM_SPLIT:
-                case TYPE_KEYGUARD:
-                case TYPE_OPTIONS_REMOTE_AND_PIP_CHANGE:
-                default:
-                    break;
-            }
         }
 
-        boolean startAnimation(
+        abstract boolean startAnimation(
                 @NonNull IBinder transition, @NonNull TransitionInfo info,
                 @NonNull SurfaceControl.Transaction startTransaction,
                 @NonNull SurfaceControl.Transaction finishTransaction,
-                @NonNull Transitions.TransitionFinishCallback finishCallback) {
-            switch (mType) {
-                case TYPE_ENTER_PIP_FROM_SPLIT:
-                    return animateEnterPipFromSplit(this, info, startTransaction, finishTransaction,
-                            finishCallback, mPlayer, mMixedHandler, mPipHandler, mSplitHandler);
-                case TYPE_ENTER_PIP_FROM_ACTIVITY_EMBEDDING:
-                    return animateEnterPipFromActivityEmbedding(
-                            info, startTransaction, finishTransaction, finishCallback);
-                case TYPE_DISPLAY_AND_SPLIT_CHANGE:
-                    return false;
-                case TYPE_OPTIONS_REMOTE_AND_PIP_CHANGE:
-                    final boolean handledToPip = animateOpenIntentWithRemoteAndPip(
-                            info, startTransaction, finishTransaction, finishCallback);
-                    // Consume the transition on remote handler if the leftover handler already
-                    // handle this transition. And if it cannot, the transition will be handled by
-                    // remote handler, so don't consume here.
-                    // Need to check leftOverHandler as it may change in
-                    // #animateOpenIntentWithRemoteAndPip
-                    if (handledToPip && mHasRequestToRemote
-                            && mLeftoversHandler != mPlayer.getRemoteTransitionHandler()) {
-                        mPlayer.getRemoteTransitionHandler().onTransitionConsumed(
-                                transition, false, null);
-                    }
-                    return handledToPip;
-                case TYPE_RECENTS_DURING_SPLIT:
-                    for (int i = info.getChanges().size() - 1; i >= 0; --i) {
-                        final TransitionInfo.Change change = info.getChanges().get(i);
-                        // Pip auto-entering info might be appended to recent transition like
-                        // pressing home-key in 3-button navigation. This offers split handler the
-                        // opportunity to handle split to pip animation.
-                        if (mPipHandler.isEnteringPip(change, info.getType())
-                                && mSplitHandler.getSplitItemPosition(change.getLastParent())
-                                != SPLIT_POSITION_UNDEFINED) {
-                            return animateEnterPipFromSplit(
-                                    this, info, startTransaction, finishTransaction, finishCallback,
-                                    mPlayer, mMixedHandler, mPipHandler, mSplitHandler);
-                        }
-                    }
+                @NonNull Transitions.TransitionFinishCallback finishCallback);
 
-                    return animateRecentsDuringSplit(
-                            info, startTransaction, finishTransaction, finishCallback);
-                case TYPE_KEYGUARD:
-                    return animateKeyguard(this, info, startTransaction, finishTransaction,
-                            finishCallback, mKeyguardHandler, mPipHandler);
-                case TYPE_RECENTS_DURING_KEYGUARD:
-                    return animateRecentsDuringKeyguard(
-                            info, startTransaction, finishTransaction, finishCallback);
-                case TYPE_RECENTS_DURING_DESKTOP:
-                    return animateRecentsDuringDesktop(
-                            info, startTransaction, finishTransaction, finishCallback);
-                case TYPE_UNFOLD:
-                    return animateUnfold(
-                            info, startTransaction, finishTransaction, finishCallback);
-                default:
-                    throw new IllegalStateException(
-                            "Starting mixed animation without a known mixed type? " + mType);
-            }
-        }
-
-        private boolean animateEnterPipFromActivityEmbedding(
-                @NonNull TransitionInfo info,
-                @NonNull SurfaceControl.Transaction startTransaction,
-                @NonNull SurfaceControl.Transaction finishTransaction,
-                @NonNull Transitions.TransitionFinishCallback finishCallback) {
-            ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, " Animating a mixed transition for "
-                    + "entering PIP from an Activity Embedding window");
-            // Split into two transitions (wct)
-            TransitionInfo.Change pipChange = null;
-            final TransitionInfo everythingElse =
-                    subCopy(info, TRANSIT_TO_BACK, true /* changes */);
-            for (int i = info.getChanges().size() - 1; i >= 0; --i) {
-                TransitionInfo.Change change = info.getChanges().get(i);
-                if (mPipHandler.isEnteringPip(change, info.getType())) {
-                    if (pipChange != null) {
-                        throw new IllegalStateException("More than 1 pip-entering changes in one"
-                                + " transition? " + info);
-                    }
-                    pipChange = change;
-                    // going backwards, so remove-by-index is fine.
-                    everythingElse.getChanges().remove(i);
-                }
-            }
-
-            final Transitions.TransitionFinishCallback finishCB = (wct) -> {
-                --mInFlightSubAnimations;
-                joinFinishArgs(wct);
-                if (mInFlightSubAnimations > 0) return;
-                finishCallback.onTransitionFinished(mFinishWCT);
-            };
-
-            if (!mActivityEmbeddingController.shouldAnimate(everythingElse)) {
-                // Fallback to dispatching to other handlers.
-                return false;
-            }
-
-            // PIP window should always be on the highest Z order.
-            if (pipChange != null) {
-                mInFlightSubAnimations = 2;
-                mPipHandler.startEnterAnimation(
-                        pipChange,
-                        startTransaction.setLayer(pipChange.getLeash(), Integer.MAX_VALUE),
-                        finishTransaction,
-                        finishCB);
-            } else {
-                mInFlightSubAnimations = 1;
-            }
-
-            mActivityEmbeddingController.startAnimation(mTransition, everythingElse,
-                    startTransaction, finishTransaction, finishCB);
-            return true;
-        }
-
-        private boolean animateOpenIntentWithRemoteAndPip(
-                @NonNull TransitionInfo info,
-                @NonNull SurfaceControl.Transaction startTransaction,
-                @NonNull SurfaceControl.Transaction finishTransaction,
-                @NonNull Transitions.TransitionFinishCallback finishCallback) {
-            TransitionInfo.Change pipChange = null;
-            for (int i = info.getChanges().size() - 1; i >= 0; --i) {
-                TransitionInfo.Change change = info.getChanges().get(i);
-                if (mPipHandler.isEnteringPip(change, info.getType())) {
-                    if (pipChange != null) {
-                        throw new IllegalStateException("More than 1 pip-entering changes in one"
-                                + " transition? " + info);
-                    }
-                    pipChange = change;
-                    info.getChanges().remove(i);
-                }
-            }
-            Transitions.TransitionFinishCallback finishCB = (wct) -> {
-                --mInFlightSubAnimations;
-                joinFinishArgs(wct);
-                if (mInFlightSubAnimations > 0) return;
-                finishCallback.onTransitionFinished(mFinishWCT);
-            };
-            if (pipChange == null) {
-                if (mLeftoversHandler != null) {
-                    mInFlightSubAnimations = 1;
-                    if (mLeftoversHandler.startAnimation(
-                            mTransition, info, startTransaction, finishTransaction, finishCB)) {
-                        return true;
-                    }
-                }
-                return false;
-            }
-            ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, " Splitting PIP into a separate"
-                    + " animation because remote-animation likely doesn't support it");
-            // Split the transition into 2 parts: the pip part and the rest.
-            mInFlightSubAnimations = 2;
-            // make a new startTransaction because pip's startEnterAnimation "consumes" it so
-            // we need a separate one to send over to launcher.
-            SurfaceControl.Transaction otherStartT = new SurfaceControl.Transaction();
-
-            mPipHandler.startEnterAnimation(pipChange, otherStartT, finishTransaction, finishCB);
-
-            // Dispatch the rest of the transition normally.
-            if (mLeftoversHandler != null
-                    && mLeftoversHandler.startAnimation(
-                            mTransition, info, startTransaction, finishTransaction, finishCB)) {
-                return true;
-            }
-            mLeftoversHandler = mPlayer.dispatchTransition(mTransition, info,
-                    startTransaction, finishTransaction, finishCB, mMixedHandler);
-            return true;
-        }
-
-        private boolean animateRecentsDuringSplit(
-                @NonNull TransitionInfo info,
-                @NonNull SurfaceControl.Transaction startTransaction,
-                @NonNull SurfaceControl.Transaction finishTransaction,
-                @NonNull Transitions.TransitionFinishCallback finishCallback) {
-            // Split-screen is only interested in the recents transition finishing (and merging), so
-            // just wrap finish and start recents animation directly.
-            Transitions.TransitionFinishCallback finishCB = (wct) -> {
-                mInFlightSubAnimations = 0;
-                // If pair-to-pair switching, the post-recents clean-up isn't needed.
-                wct = wct != null ? wct : new WindowContainerTransaction();
-                if (mAnimType != ANIM_TYPE_PAIR_TO_PAIR) {
-                    mSplitHandler.onRecentsInSplitAnimationFinish(wct, finishTransaction);
-                } else {
-                    // notify pair-to-pair recents animation finish
-                    mSplitHandler.onRecentsPairToPairAnimationFinish(wct);
-                }
-                mSplitHandler.onTransitionAnimationComplete();
-                finishCallback.onTransitionFinished(wct);
-            };
-            mInFlightSubAnimations = 1;
-            mSplitHandler.onRecentsInSplitAnimationStart(info);
-            final boolean handled = mLeftoversHandler.startAnimation(mTransition, info,
-                    startTransaction, finishTransaction, finishCB);
-            if (!handled) {
-                mSplitHandler.onRecentsInSplitAnimationCanceled();
-            }
-            return handled;
-        }
-
-        private boolean animateRecentsDuringKeyguard(
-                @NonNull TransitionInfo info,
-                @NonNull SurfaceControl.Transaction startTransaction,
-                @NonNull SurfaceControl.Transaction finishTransaction,
-                @NonNull Transitions.TransitionFinishCallback finishCallback) {
-            if (mInfo == null) {
-                mInfo = info;
-                mFinishT = finishTransaction;
-                mFinishCB = finishCallback;
-            }
-            return startSubAnimation(mRecentsHandler, info, startTransaction, finishTransaction);
-        }
-
-        private boolean animateRecentsDuringDesktop(
-                @NonNull TransitionInfo info,
-                @NonNull SurfaceControl.Transaction startTransaction,
-                @NonNull SurfaceControl.Transaction finishTransaction,
-                @NonNull Transitions.TransitionFinishCallback finishCallback) {
-            Transitions.TransitionFinishCallback finishCB = wct -> {
-                mInFlightSubAnimations--;
-                if (mInFlightSubAnimations == 0) {
-                    finishCallback.onTransitionFinished(wct);
-                }
-            };
-
-            mInFlightSubAnimations++;
-            boolean consumed = mRecentsHandler.startAnimation(
-                    mTransition, info, startTransaction, finishTransaction, finishCB);
-            if (!consumed) {
-                mInFlightSubAnimations--;
-                return false;
-            }
-            if (mDesktopTasksController != null) {
-                mDesktopTasksController.syncSurfaceState(info, finishTransaction);
-                return true;
-            }
-
-            return false;
-        }
-
-        private boolean animateUnfold(
-                @NonNull TransitionInfo info,
-                @NonNull SurfaceControl.Transaction startTransaction,
-                @NonNull SurfaceControl.Transaction finishTransaction,
-                @NonNull Transitions.TransitionFinishCallback finishCallback) {
-            final Transitions.TransitionFinishCallback finishCB = (wct) -> {
-                mInFlightSubAnimations--;
-                if (mInFlightSubAnimations > 0) return;
-                finishCallback.onTransitionFinished(wct);
-            };
-            mInFlightSubAnimations = 1;
-            // Sync pip state.
-            if (mPipHandler != null) {
-                mPipHandler.syncPipSurfaceState(info, startTransaction, finishTransaction);
-            }
-            if (mSplitHandler != null && mSplitHandler.isSplitActive()) {
-                mSplitHandler.updateSurfaces(startTransaction);
-            }
-            return mUnfoldHandler.startAnimation(
-                    mTransition, info, startTransaction, finishTransaction, finishCB);
-        }
-
-        void mergeAnimation(
+        abstract void mergeAnimation(
                 @NonNull IBinder transition, @NonNull TransitionInfo info,
                 @NonNull SurfaceControl.Transaction t, @NonNull IBinder mergeTarget,
-                @NonNull Transitions.TransitionFinishCallback finishCallback) {
-            switch (mType) {
-                case TYPE_DISPLAY_AND_SPLIT_CHANGE:
-                    // queue since no actual animation.
-                    break;
-                case TYPE_ENTER_PIP_FROM_SPLIT:
-                    if (mAnimType == ANIM_TYPE_GOING_HOME) {
-                        boolean ended = mSplitHandler.end();
-                        // If split couldn't end (because it is remote), then don't end everything
-                        // else since we have to play out the animation anyways.
-                        if (!ended) return;
-                        mPipHandler.end();
-                        if (mLeftoversHandler != null) {
-                            mLeftoversHandler.mergeAnimation(
-                                    transition, info, t, mergeTarget, finishCallback);
-                        }
-                    } else {
-                        mPipHandler.end();
-                    }
-                    break;
-                case TYPE_ENTER_PIP_FROM_ACTIVITY_EMBEDDING:
-                    mPipHandler.end();
-                    mActivityEmbeddingController.mergeAnimation(transition, info, t, mergeTarget,
-                            finishCallback);
-                    break;
-                case TYPE_OPTIONS_REMOTE_AND_PIP_CHANGE:
-                    mPipHandler.end();
-                    if (mLeftoversHandler != null) {
-                        mLeftoversHandler.mergeAnimation(transition, info, t, mergeTarget,
-                                finishCallback);
-                    }
-                    break;
-                case TYPE_RECENTS_DURING_SPLIT:
-                    if (mSplitHandler.isPendingEnter(transition)) {
-                        // Recents -> enter-split means that we are switching from one pair to
-                        // another pair.
-                        mAnimType = ANIM_TYPE_PAIR_TO_PAIR;
-                    }
-                    mLeftoversHandler.mergeAnimation(
-                            transition, info, t, mergeTarget, finishCallback);
-                    break;
-                case TYPE_KEYGUARD:
-                    mKeyguardHandler.mergeAnimation(
-                            transition, info, t, mergeTarget, finishCallback);
-                    break;
-                case TYPE_RECENTS_DURING_KEYGUARD:
-                    if ((info.getFlags() & TRANSIT_FLAG_KEYGUARD_UNOCCLUDING) != 0) {
-                        DefaultMixedHandler.handoverTransitionLeashes(mInfo, info, t, mFinishT);
-                        if (animateKeyguard(this, info, t, mFinishT, mFinishCB, mKeyguardHandler,
-                                mPipHandler)) {
-                            finishCallback.onTransitionFinished(null);
-                        }
-                    }
-                    mLeftoversHandler.mergeAnimation(
-                            transition, info, t, mergeTarget, finishCallback);
-                    break;
-                case TYPE_RECENTS_DURING_DESKTOP:
-                    mLeftoversHandler.mergeAnimation(
-                            transition, info, t, mergeTarget, finishCallback);
-                    break;
-                case TYPE_UNFOLD:
-                    mUnfoldHandler.mergeAnimation(transition, info, t, mergeTarget, finishCallback);
-                    break;
-                default:
-                    throw new IllegalStateException(
-                            "Playing a mixed transition with unknown type? " + mType);
-            }
-        }
+                @NonNull Transitions.TransitionFinishCallback finishCallback);
 
-        void onTransitionConsumed(
+        abstract void onTransitionConsumed(
                 @NonNull IBinder transition, boolean aborted,
-                @Nullable SurfaceControl.Transaction finishT) {
-            switch (mType) {
-                case TYPE_ENTER_PIP_FROM_SPLIT:
-                    mPipHandler.onTransitionConsumed(transition, aborted, finishT);
-                    break;
-                case TYPE_ENTER_PIP_FROM_ACTIVITY_EMBEDDING:
-                    mPipHandler.onTransitionConsumed(transition, aborted, finishT);
-                    mActivityEmbeddingController.onTransitionConsumed(transition, aborted, finishT);
-                    break;
-                case TYPE_RECENTS_DURING_SPLIT:
-                case TYPE_OPTIONS_REMOTE_AND_PIP_CHANGE:
-                case TYPE_RECENTS_DURING_DESKTOP:
-                    mLeftoversHandler.onTransitionConsumed(transition, aborted, finishT);
-                    break;
-                case TYPE_KEYGUARD:
-                    mKeyguardHandler.onTransitionConsumed(transition, aborted, finishT);
-                    break;
-                case TYPE_UNFOLD:
-                    mUnfoldHandler.onTransitionConsumed(transition, aborted, finishT);
-                    break;
-                default:
-                    break;
-            }
+                @Nullable SurfaceControl.Transaction finishT);
 
-            if (mHasRequestToRemote) {
-                mPlayer.getRemoteTransitionHandler().onTransitionConsumed(
-                        transition, aborted, finishT);
-            }
-        }
-
-        boolean startSubAnimation(Transitions.TransitionHandler handler, TransitionInfo info,
+        protected boolean startSubAnimation(
+                Transitions.TransitionHandler handler, TransitionInfo info,
                 SurfaceControl.Transaction startT, SurfaceControl.Transaction finishT) {
             if (mInfo != null) {
                 ProtoLog.d(ShellProtoLogGroup.WM_SHELL_TRANSITIONS,
@@ -573,7 +184,7 @@
             return true;
         }
 
-        void onSubAnimationFinished(TransitionInfo info, WindowContainerTransaction wct) {
+        private void onSubAnimationFinished(TransitionInfo info, WindowContainerTransaction wct) {
             mInFlightSubAnimations--;
             if (mInfo != null) {
                 ProtoLog.d(ShellProtoLogGroup.WM_SHELL_TRANSITIONS,
@@ -644,7 +255,7 @@
                 throw new IllegalStateException("Unexpected remote transition in"
                         + "pip-enter-from-split request");
             }
-            mActiveTransitions.add(createMixedTransition(
+            mActiveTransitions.add(createDefaultMixedTransition(
                     MixedTransition.TYPE_ENTER_PIP_FROM_SPLIT, transition));
 
             WindowContainerTransaction out = new WindowContainerTransaction();
@@ -656,7 +267,7 @@
                 mActivityEmbeddingController != null)) {
             ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS,
                     " Got a PiP-enter request from an Activity Embedding split");
-            mActiveTransitions.add(createMixedTransition(
+            mActiveTransitions.add(createDefaultMixedTransition(
                     MixedTransition.TYPE_ENTER_PIP_FROM_ACTIVITY_EMBEDDING, transition));
             // Postpone transition splitting to later.
             WindowContainerTransaction out = new WindowContainerTransaction();
@@ -675,7 +286,7 @@
             if (handler == null) {
                 return null;
             }
-            final MixedTransition mixed = createMixedTransition(
+            final MixedTransition mixed = createDefaultMixedTransition(
                     MixedTransition.TYPE_OPTIONS_REMOTE_AND_PIP_CHANGE, transition);
             mixed.mLeftoversHandler = handler.first;
             mActiveTransitions.add(mixed);
@@ -701,7 +312,7 @@
                         mPlayer.getRemoteTransitionHandler(),
                         new WindowContainerTransaction());
             }
-            final MixedTransition mixed = createMixedTransition(
+            final MixedTransition mixed = createRecentsMixedTransition(
                     MixedTransition.TYPE_RECENTS_DURING_SPLIT, transition);
             mixed.mLeftoversHandler = handler.first;
             mActiveTransitions.add(mixed);
@@ -710,7 +321,7 @@
             final WindowContainerTransaction wct =
                     mUnfoldHandler.handleRequest(transition, request);
             if (wct != null) {
-                mActiveTransitions.add(createMixedTransition(
+                mActiveTransitions.add(createDefaultMixedTransition(
                         MixedTransition.TYPE_UNFOLD, transition));
             }
             return wct;
@@ -718,6 +329,12 @@
         return null;
     }
 
+    private DefaultMixedTransition createDefaultMixedTransition(int type, IBinder transition) {
+        return new DefaultMixedTransition(
+                type, transition, mPlayer, this, mPipHandler, mSplitHandler, mKeyguardHandler,
+                mUnfoldHandler, mActivityEmbeddingController);
+    }
+
     @Override
     public Consumer<IBinder> handleRecentsRequest(WindowContainerTransaction outWCT) {
         if (mRecentsHandler != null) {
@@ -737,31 +354,30 @@
     private void setRecentsTransitionDuringSplit(IBinder transition) {
         ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, " Got a recents request while "
                 + "Split-Screen is foreground, so treat it as Mixed.");
-        mActiveTransitions.add(createMixedTransition(
+        mActiveTransitions.add(createRecentsMixedTransition(
                 MixedTransition.TYPE_RECENTS_DURING_SPLIT, transition));
     }
 
     private void setRecentsTransitionDuringKeyguard(IBinder transition) {
         ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, " Got a recents request while "
                 + "keyguard is visible, so treat it as Mixed.");
-        mActiveTransitions.add(createMixedTransition(
+        mActiveTransitions.add(createRecentsMixedTransition(
                 MixedTransition.TYPE_RECENTS_DURING_KEYGUARD, transition));
     }
 
     private void setRecentsTransitionDuringDesktop(IBinder transition) {
         ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, " Got a recents request while "
                 + "desktop mode is active, so treat it as Mixed.");
-        mActiveTransitions.add(createMixedTransition(
+        mActiveTransitions.add(createRecentsMixedTransition(
                 MixedTransition.TYPE_RECENTS_DURING_DESKTOP, transition));
     }
 
-    private MixedTransition createMixedTransition(int type, IBinder transition) {
-        return new MixedTransition(type, transition, mPlayer, this, mPipHandler, mRecentsHandler,
-                mSplitHandler, mKeyguardHandler, mDesktopTasksController, mUnfoldHandler,
-                mActivityEmbeddingController);
+    private MixedTransition createRecentsMixedTransition(int type, IBinder transition) {
+        return new RecentsMixedTransition(type, transition, mPlayer, this, mPipHandler,
+                mSplitHandler, mKeyguardHandler, mRecentsHandler, mDesktopTasksController);
     }
 
-    private static TransitionInfo subCopy(@NonNull TransitionInfo info,
+    static TransitionInfo subCopy(@NonNull TransitionInfo info,
             @WindowManager.TransitionType int newType, boolean withChanges) {
         final TransitionInfo out = new TransitionInfo(newType, withChanges ? info.getFlags() : 0);
         out.setTrack(info.getTrack());
@@ -778,15 +394,6 @@
         return out;
     }
 
-    private static boolean isHomeOpening(@NonNull TransitionInfo.Change change) {
-        return change.getTaskInfo() != null
-                && change.getTaskInfo().getActivityType() == ACTIVITY_TYPE_HOME;
-    }
-
-    private static boolean isWallpaper(@NonNull TransitionInfo.Change change) {
-        return (change.getFlags() & FLAG_IS_WALLPAPER) != 0;
-    }
-
     @Override
     public boolean startAnimation(@NonNull IBinder transition, @NonNull TransitionInfo info,
             @NonNull SurfaceControl.Transaction startTransaction,
@@ -805,7 +412,7 @@
         if (KeyguardTransitionHandler.handles(info)) {
             if (mixed != null && mixed.mType != MixedTransition.TYPE_KEYGUARD) {
                 final MixedTransition keyguardMixed =
-                        createMixedTransition(MixedTransition.TYPE_KEYGUARD, transition);
+                        createDefaultMixedTransition(MixedTransition.TYPE_KEYGUARD, transition);
                 mActiveTransitions.add(keyguardMixed);
                 Transitions.TransitionFinishCallback callback = wct -> {
                     mActiveTransitions.remove(keyguardMixed);
@@ -845,117 +452,6 @@
         return handled;
     }
 
-    private static boolean animateEnterPipFromSplit(@NonNull final MixedTransition mixed,
-            @NonNull TransitionInfo info,
-            @NonNull SurfaceControl.Transaction startTransaction,
-            @NonNull SurfaceControl.Transaction finishTransaction,
-            @NonNull Transitions.TransitionFinishCallback finishCallback,
-            @NonNull Transitions player, @NonNull DefaultMixedHandler mixedHandler,
-            @NonNull PipTransitionController pipHandler, @NonNull StageCoordinator splitHandler) {
-        ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, " Animating a mixed transition for "
-                + "entering PIP while Split-Screen is foreground.");
-        TransitionInfo.Change pipChange = null;
-        TransitionInfo.Change wallpaper = null;
-        final TransitionInfo everythingElse = subCopy(info, TRANSIT_TO_BACK, true /* changes */);
-        boolean homeIsOpening = false;
-        for (int i = info.getChanges().size() - 1; i >= 0; --i) {
-            TransitionInfo.Change change = info.getChanges().get(i);
-            if (pipHandler.isEnteringPip(change, info.getType())) {
-                if (pipChange != null) {
-                    throw new IllegalStateException("More than 1 pip-entering changes in one"
-                            + " transition? " + info);
-                }
-                pipChange = change;
-                // going backwards, so remove-by-index is fine.
-                everythingElse.getChanges().remove(i);
-            } else if (isHomeOpening(change)) {
-                homeIsOpening = true;
-            } else if (isWallpaper(change)) {
-                wallpaper = change;
-            }
-        }
-        if (pipChange == null) {
-            // um, something probably went wrong.
-            return false;
-        }
-        final boolean isGoingHome = homeIsOpening;
-        Transitions.TransitionFinishCallback finishCB = (wct) -> {
-            --mixed.mInFlightSubAnimations;
-            mixed.joinFinishArgs(wct);
-            if (mixed.mInFlightSubAnimations > 0) return;
-            if (isGoingHome) {
-                splitHandler.onTransitionAnimationComplete();
-            }
-            finishCallback.onTransitionFinished(mixed.mFinishWCT);
-        };
-        if (isGoingHome || splitHandler.getSplitItemPosition(pipChange.getLastParent())
-                != SPLIT_POSITION_UNDEFINED) {
-            ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, " Animation is actually mixed "
-                    + "since entering-PiP caused us to leave split and return home.");
-            // We need to split the transition into 2 parts: the pip part (animated by pip)
-            // and the dismiss-part (animated by launcher).
-            mixed.mInFlightSubAnimations = 2;
-            // immediately make the wallpaper visible (so that we don't see it pop-in during
-            // the time it takes to start recents animation (which is remote).
-            if (wallpaper != null) {
-                startTransaction.show(wallpaper.getLeash()).setAlpha(wallpaper.getLeash(), 1.f);
-            }
-            // make a new startTransaction because pip's startEnterAnimation "consumes" it so
-            // we need a separate one to send over to launcher.
-            SurfaceControl.Transaction otherStartT = new SurfaceControl.Transaction();
-            @SplitScreen.StageType int topStageToKeep = STAGE_TYPE_UNDEFINED;
-            if (splitHandler.isSplitScreenVisible()) {
-                // The non-going home case, we could be pip-ing one of the split stages and keep
-                // showing the other
-                for (int i = info.getChanges().size() - 1; i >= 0; --i) {
-                    TransitionInfo.Change change = info.getChanges().get(i);
-                    if (change == pipChange) {
-                        // Ignore the change/task that's going into Pip
-                        continue;
-                    }
-                    @SplitScreen.StageType int splitItemStage =
-                            splitHandler.getSplitItemStage(change.getLastParent());
-                    if (splitItemStage != STAGE_TYPE_UNDEFINED) {
-                        topStageToKeep = splitItemStage;
-                        break;
-                    }
-                }
-            }
-            // Let split update internal state for dismiss.
-            splitHandler.prepareDismissAnimation(topStageToKeep,
-                    EXIT_REASON_CHILD_TASK_ENTER_PIP, everythingElse, otherStartT,
-                    finishTransaction);
-
-            // We are trying to accommodate launcher's close animation which can't handle the
-            // divider-bar, so if split-handler is closing the divider-bar, just hide it and remove
-            // from transition info.
-            for (int i = everythingElse.getChanges().size() - 1; i >= 0; --i) {
-                if ((everythingElse.getChanges().get(i).getFlags() & FLAG_IS_DIVIDER_BAR) != 0) {
-                    everythingElse.getChanges().remove(i);
-                    break;
-                }
-            }
-
-            pipHandler.setEnterAnimationType(ANIM_TYPE_ALPHA);
-            pipHandler.startEnterAnimation(pipChange, startTransaction, finishTransaction,
-                    finishCB);
-            // Dispatch the rest of the transition normally. This will most-likely be taken by
-            // recents or default handler.
-            mixed.mLeftoversHandler = player.dispatchTransition(mixed.mTransition, everythingElse,
-                    otherStartT, finishTransaction, finishCB, mixedHandler);
-        } else {
-            ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, "  Not leaving split, so just "
-                    + "forward animation to Pip-Handler.");
-            // This happens if the pip-ing activity is in a multi-activity task (and thus a
-            // new pip task is spawned). In this case, we don't actually exit split so we can
-            // just let pip transition handle the animation verbatim.
-            mixed.mInFlightSubAnimations = 1;
-            pipHandler.startAnimation(mixed.mTransition, info, startTransaction, finishTransaction,
-                    finishCB);
-        }
-        return true;
-    }
-
     private void unlinkMissingParents(TransitionInfo from) {
         for (int i = 0; i < from.getChanges().size(); ++i) {
             final TransitionInfo.Change chg = from.getChanges().get(i);
@@ -987,15 +483,14 @@
     public boolean animatePendingEnterPipFromSplit(IBinder transition, TransitionInfo info,
             SurfaceControl.Transaction startT, SurfaceControl.Transaction finishT,
             Transitions.TransitionFinishCallback finishCallback) {
-        final MixedTransition mixed = createMixedTransition(
+        final MixedTransition mixed = createDefaultMixedTransition(
                 MixedTransition.TYPE_ENTER_PIP_FROM_SPLIT, transition);
         mActiveTransitions.add(mixed);
         Transitions.TransitionFinishCallback callback = wct -> {
             mActiveTransitions.remove(mixed);
             finishCallback.onTransitionFinished(wct);
         };
-        return animateEnterPipFromSplit(mixed, info, startT, finishT, finishCallback, mPlayer, this,
-                mPipHandler, mSplitHandler);
+        return mixed.startAnimation(transition, info, startT, finishT, callback);
     }
 
     /**
@@ -1018,7 +513,7 @@
         }
         if (displayPart.getChanges().isEmpty()) return false;
         unlinkMissingParents(everythingElse);
-        final MixedTransition mixed = createMixedTransition(
+        final MixedTransition mixed = createDefaultMixedTransition(
                 MixedTransition.TYPE_DISPLAY_AND_SPLIT_CHANGE, transition);
         mActiveTransitions.add(mixed);
         ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, " Animation is a mix of display change "
@@ -1135,7 +630,7 @@
      * {@link TransitionInfo} so that it can take over some parts of the animation without
      * reparenting to new transition roots.
      */
-    private static void handoverTransitionLeashes(
+    static void handoverTransitionLeashes(
             @NonNull TransitionInfo from,
             @NonNull TransitionInfo to,
             @NonNull SurfaceControl.Transaction startT,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedTransition.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedTransition.java
new file mode 100644
index 0000000..9ce46d6
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedTransition.java
@@ -0,0 +1,315 @@
+/*
+ * 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.transition;
+
+import static android.view.WindowManager.TRANSIT_TO_BACK;
+
+import static com.android.wm.shell.transition.DefaultMixedHandler.subCopy;
+import static com.android.wm.shell.transition.MixedTransitionHelper.animateEnterPipFromSplit;
+import static com.android.wm.shell.transition.MixedTransitionHelper.animateKeyguard;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.os.IBinder;
+import android.view.SurfaceControl;
+import android.window.TransitionInfo;
+
+import com.android.internal.protolog.common.ProtoLog;
+import com.android.wm.shell.activityembedding.ActivityEmbeddingController;
+import com.android.wm.shell.keyguard.KeyguardTransitionHandler;
+import com.android.wm.shell.pip.PipTransitionController;
+import com.android.wm.shell.protolog.ShellProtoLogGroup;
+import com.android.wm.shell.splitscreen.StageCoordinator;
+import com.android.wm.shell.unfold.UnfoldTransitionHandler;
+
+class DefaultMixedTransition extends DefaultMixedHandler.MixedTransition {
+    private final UnfoldTransitionHandler mUnfoldHandler;
+    private final ActivityEmbeddingController mActivityEmbeddingController;
+
+    DefaultMixedTransition(int type, IBinder transition, Transitions player,
+            DefaultMixedHandler mixedHandler, PipTransitionController pipHandler,
+            StageCoordinator splitHandler, KeyguardTransitionHandler keyguardHandler,
+            UnfoldTransitionHandler unfoldHandler,
+            ActivityEmbeddingController activityEmbeddingController) {
+        super(type, transition, player, mixedHandler, pipHandler, splitHandler, keyguardHandler);
+        mUnfoldHandler = unfoldHandler;
+        mActivityEmbeddingController = activityEmbeddingController;
+
+        switch (type) {
+            case TYPE_UNFOLD:
+                mLeftoversHandler = mUnfoldHandler;
+                break;
+            case TYPE_DISPLAY_AND_SPLIT_CHANGE:
+            case TYPE_ENTER_PIP_FROM_ACTIVITY_EMBEDDING:
+            case TYPE_ENTER_PIP_FROM_SPLIT:
+            case TYPE_KEYGUARD:
+            case TYPE_OPTIONS_REMOTE_AND_PIP_CHANGE:
+            default:
+                break;
+        }
+    }
+
+    @Override
+    boolean startAnimation(
+            @NonNull IBinder transition, @NonNull TransitionInfo info,
+            @NonNull SurfaceControl.Transaction startTransaction,
+            @NonNull SurfaceControl.Transaction finishTransaction,
+            @NonNull Transitions.TransitionFinishCallback finishCallback) {
+        return switch (mType) {
+            case TYPE_DISPLAY_AND_SPLIT_CHANGE -> false;
+            case TYPE_ENTER_PIP_FROM_ACTIVITY_EMBEDDING ->
+                    animateEnterPipFromActivityEmbedding(
+                            info, startTransaction, finishTransaction, finishCallback);
+            case TYPE_ENTER_PIP_FROM_SPLIT ->
+                    animateEnterPipFromSplit(this, info, startTransaction, finishTransaction,
+                            finishCallback, mPlayer, mMixedHandler, mPipHandler, mSplitHandler);
+            case TYPE_KEYGUARD ->
+                    animateKeyguard(this, info, startTransaction, finishTransaction, finishCallback,
+                            mKeyguardHandler, mPipHandler);
+            case TYPE_OPTIONS_REMOTE_AND_PIP_CHANGE ->
+                    animateOpenIntentWithRemoteAndPip(transition, info, startTransaction,
+                            finishTransaction, finishCallback);
+            case TYPE_UNFOLD ->
+                    animateUnfold(info, startTransaction, finishTransaction, finishCallback);
+            default -> throw new IllegalStateException(
+                    "Starting default mixed animation with unknown or illegal type: " + mType);
+        };
+    }
+
+    private boolean animateEnterPipFromActivityEmbedding(
+            @NonNull TransitionInfo info,
+            @NonNull SurfaceControl.Transaction startTransaction,
+            @NonNull SurfaceControl.Transaction finishTransaction,
+            @NonNull Transitions.TransitionFinishCallback finishCallback) {
+        ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, " Animating a mixed transition for "
+                + "entering PIP from an Activity Embedding window");
+        // Split into two transitions (wct)
+        TransitionInfo.Change pipChange = null;
+        final TransitionInfo everythingElse = subCopy(info, TRANSIT_TO_BACK, true /* changes */);
+        for (int i = info.getChanges().size() - 1; i >= 0; --i) {
+            TransitionInfo.Change change = info.getChanges().get(i);
+            if (mPipHandler.isEnteringPip(change, info.getType())) {
+                if (pipChange != null) {
+                    throw new IllegalStateException("More than 1 pip-entering changes in one"
+                            + " transition? " + info);
+                }
+                pipChange = change;
+                // going backwards, so remove-by-index is fine.
+                everythingElse.getChanges().remove(i);
+            }
+        }
+
+        final Transitions.TransitionFinishCallback finishCB = (wct) -> {
+            --mInFlightSubAnimations;
+            joinFinishArgs(wct);
+            if (mInFlightSubAnimations > 0) return;
+            finishCallback.onTransitionFinished(mFinishWCT);
+        };
+
+        if (!mActivityEmbeddingController.shouldAnimate(everythingElse)) {
+            // Fallback to dispatching to other handlers.
+            return false;
+        }
+
+        // PIP window should always be on the highest Z order.
+        if (pipChange != null) {
+            mInFlightSubAnimations = 2;
+            mPipHandler.startEnterAnimation(
+                    pipChange, startTransaction.setLayer(pipChange.getLeash(), Integer.MAX_VALUE),
+                    finishTransaction,
+                    finishCB);
+        } else {
+            mInFlightSubAnimations = 1;
+        }
+
+        mActivityEmbeddingController.startAnimation(
+                mTransition, everythingElse, startTransaction, finishTransaction, finishCB);
+        return true;
+    }
+
+    private boolean animateOpenIntentWithRemoteAndPip(
+            @NonNull IBinder transition, @NonNull TransitionInfo info,
+            @NonNull SurfaceControl.Transaction startTransaction,
+            @NonNull SurfaceControl.Transaction finishTransaction,
+            @NonNull Transitions.TransitionFinishCallback finishCallback) {
+        boolean handledToPip = tryAnimateOpenIntentWithRemoteAndPip(
+                info, startTransaction, finishTransaction, finishCallback);
+        // Consume the transition on remote handler if the leftover handler already handle this
+        // transition. And if it cannot, the transition will be handled by remote handler, so don't
+        // consume here.
+        // Need to check leftOverHandler as it may change in #animateOpenIntentWithRemoteAndPip
+        if (handledToPip && mHasRequestToRemote
+                && mLeftoversHandler != mPlayer.getRemoteTransitionHandler()) {
+            mPlayer.getRemoteTransitionHandler().onTransitionConsumed(transition, false, null);
+        }
+        return handledToPip;
+    }
+
+    private boolean tryAnimateOpenIntentWithRemoteAndPip(
+            @NonNull TransitionInfo info,
+            @NonNull SurfaceControl.Transaction startTransaction,
+            @NonNull SurfaceControl.Transaction finishTransaction,
+            @NonNull Transitions.TransitionFinishCallback finishCallback) {
+        TransitionInfo.Change pipChange = null;
+        for (int i = info.getChanges().size() - 1; i >= 0; --i) {
+            TransitionInfo.Change change = info.getChanges().get(i);
+            if (mPipHandler.isEnteringPip(change, info.getType())) {
+                if (pipChange != null) {
+                    throw new IllegalStateException("More than 1 pip-entering changes in one"
+                            + " transition? " + info);
+                }
+                pipChange = change;
+                info.getChanges().remove(i);
+            }
+        }
+        Transitions.TransitionFinishCallback finishCB = (wct) -> {
+            --mInFlightSubAnimations;
+            joinFinishArgs(wct);
+            if (mInFlightSubAnimations > 0) return;
+            finishCallback.onTransitionFinished(mFinishWCT);
+        };
+        if (pipChange == null) {
+            if (mLeftoversHandler != null) {
+                mInFlightSubAnimations = 1;
+                if (mLeftoversHandler.startAnimation(
+                        mTransition, info, startTransaction, finishTransaction, finishCB)) {
+                    return true;
+                }
+            }
+            return false;
+        }
+        ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, " Splitting PIP into a separate"
+                + " animation because remote-animation likely doesn't support it");
+        // Split the transition into 2 parts: the pip part and the rest.
+        mInFlightSubAnimations = 2;
+        // make a new startTransaction because pip's startEnterAnimation "consumes" it so
+        // we need a separate one to send over to launcher.
+        SurfaceControl.Transaction otherStartT = new SurfaceControl.Transaction();
+
+        mPipHandler.startEnterAnimation(pipChange, otherStartT, finishTransaction, finishCB);
+
+        // Dispatch the rest of the transition normally.
+        if (mLeftoversHandler != null
+                && mLeftoversHandler.startAnimation(mTransition, info,
+                startTransaction, finishTransaction, finishCB)) {
+            return true;
+        }
+        mLeftoversHandler = mPlayer.dispatchTransition(
+                mTransition, info, startTransaction, finishTransaction, finishCB, mMixedHandler);
+        return true;
+    }
+
+    private boolean animateUnfold(
+            @NonNull TransitionInfo info,
+            @NonNull SurfaceControl.Transaction startTransaction,
+            @NonNull SurfaceControl.Transaction finishTransaction,
+            @NonNull Transitions.TransitionFinishCallback finishCallback) {
+        final Transitions.TransitionFinishCallback finishCB = (wct) -> {
+            mInFlightSubAnimations--;
+            if (mInFlightSubAnimations > 0) return;
+            finishCallback.onTransitionFinished(wct);
+        };
+        mInFlightSubAnimations = 1;
+        // Sync pip state.
+        if (mPipHandler != null) {
+            mPipHandler.syncPipSurfaceState(info, startTransaction, finishTransaction);
+        }
+        if (mSplitHandler != null && mSplitHandler.isSplitActive()) {
+            mSplitHandler.updateSurfaces(startTransaction);
+        }
+        return mUnfoldHandler.startAnimation(
+                mTransition, info, startTransaction, finishTransaction, finishCB);
+    }
+
+    @Override
+    void mergeAnimation(
+            @NonNull IBinder transition, @NonNull TransitionInfo info,
+            @NonNull SurfaceControl.Transaction t, @NonNull IBinder mergeTarget,
+            @NonNull Transitions.TransitionFinishCallback finishCallback) {
+        switch (mType) {
+            case TYPE_DISPLAY_AND_SPLIT_CHANGE:
+                // queue since no actual animation.
+                return;
+            case TYPE_ENTER_PIP_FROM_ACTIVITY_EMBEDDING:
+                mPipHandler.end();
+                mActivityEmbeddingController.mergeAnimation(
+                        transition, info, t, mergeTarget, finishCallback);
+                return;
+            case TYPE_ENTER_PIP_FROM_SPLIT:
+                if (mAnimType == ANIM_TYPE_GOING_HOME) {
+                    boolean ended = mSplitHandler.end();
+                    // If split couldn't end (because it is remote), then don't end everything else
+                    // since we have to play out the animation anyways.
+                    if (!ended) return;
+                    mPipHandler.end();
+                    if (mLeftoversHandler != null) {
+                        mLeftoversHandler.mergeAnimation(
+                                transition, info, t, mergeTarget, finishCallback);
+                    }
+                } else {
+                    mPipHandler.end();
+                }
+                return;
+            case TYPE_KEYGUARD:
+                mKeyguardHandler.mergeAnimation(transition, info, t, mergeTarget, finishCallback);
+                return;
+            case TYPE_OPTIONS_REMOTE_AND_PIP_CHANGE:
+                mPipHandler.end();
+                if (mLeftoversHandler != null) {
+                    mLeftoversHandler.mergeAnimation(
+                            transition, info, t, mergeTarget, finishCallback);
+                }
+                return;
+            case TYPE_UNFOLD:
+                mUnfoldHandler.mergeAnimation(transition, info, t, mergeTarget, finishCallback);
+                return;
+            default:
+                throw new IllegalStateException("Playing a default mixed transition with unknown or"
+                        + " illegal type: " + mType);
+        }
+    }
+
+    @Override
+    void onTransitionConsumed(
+            @NonNull IBinder transition, boolean aborted,
+            @Nullable SurfaceControl.Transaction finishT) {
+        switch (mType) {
+            case TYPE_ENTER_PIP_FROM_ACTIVITY_EMBEDDING:
+                mPipHandler.onTransitionConsumed(transition, aborted, finishT);
+                mActivityEmbeddingController.onTransitionConsumed(transition, aborted, finishT);
+                break;
+            case TYPE_ENTER_PIP_FROM_SPLIT:
+                mPipHandler.onTransitionConsumed(transition, aborted, finishT);
+                break;
+            case TYPE_KEYGUARD:
+                mKeyguardHandler.onTransitionConsumed(transition, aborted, finishT);
+                break;
+            case TYPE_OPTIONS_REMOTE_AND_PIP_CHANGE:
+                mLeftoversHandler.onTransitionConsumed(transition, aborted, finishT);
+                break;
+            case TYPE_UNFOLD:
+                mUnfoldHandler.onTransitionConsumed(transition, aborted, finishT);
+                break;
+            default:
+                break;
+        }
+
+        if (mHasRequestToRemote) {
+            mPlayer.getRemoteTransitionHandler().onTransitionConsumed(transition, aborted, finishT);
+        }
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/MixedTransitionHelper.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/MixedTransitionHelper.java
new file mode 100644
index 0000000..0974cd1
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/MixedTransitionHelper.java
@@ -0,0 +1,181 @@
+/*
+ * 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.transition;
+
+import static android.app.WindowConfiguration.ACTIVITY_TYPE_HOME;
+import static android.view.WindowManager.TRANSIT_TO_BACK;
+import static android.window.TransitionInfo.FLAG_IS_WALLPAPER;
+
+import static com.android.wm.shell.common.split.SplitScreenConstants.FLAG_IS_DIVIDER_BAR;
+import static com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_UNDEFINED;
+import static com.android.wm.shell.pip.PipAnimationController.ANIM_TYPE_ALPHA;
+import static com.android.wm.shell.splitscreen.SplitScreen.STAGE_TYPE_UNDEFINED;
+import static com.android.wm.shell.splitscreen.SplitScreenController.EXIT_REASON_CHILD_TASK_ENTER_PIP;
+import static com.android.wm.shell.transition.DefaultMixedHandler.subCopy;
+
+import android.annotation.NonNull;
+import android.view.SurfaceControl;
+import android.window.TransitionInfo;
+
+import com.android.internal.protolog.common.ProtoLog;
+import com.android.wm.shell.keyguard.KeyguardTransitionHandler;
+import com.android.wm.shell.pip.PipTransitionController;
+import com.android.wm.shell.protolog.ShellProtoLogGroup;
+import com.android.wm.shell.splitscreen.SplitScreen;
+import com.android.wm.shell.splitscreen.StageCoordinator;
+
+public class MixedTransitionHelper {
+    static boolean animateEnterPipFromSplit(
+            @NonNull DefaultMixedHandler.MixedTransition mixed, @NonNull TransitionInfo info,
+            @NonNull SurfaceControl.Transaction startTransaction,
+            @NonNull SurfaceControl.Transaction finishTransaction,
+            @NonNull Transitions.TransitionFinishCallback finishCallback,
+            @NonNull Transitions player, @NonNull DefaultMixedHandler mixedHandler,
+            @NonNull PipTransitionController pipHandler, @NonNull StageCoordinator splitHandler) {
+        ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, " Animating a mixed transition for "
+                + "entering PIP while Split-Screen is foreground.");
+        TransitionInfo.Change pipChange = null;
+        TransitionInfo.Change wallpaper = null;
+        final TransitionInfo everythingElse =
+                subCopy(info, TRANSIT_TO_BACK, true /* changes */);
+        boolean homeIsOpening = false;
+        for (int i = info.getChanges().size() - 1; i >= 0; --i) {
+            TransitionInfo.Change change = info.getChanges().get(i);
+            if (pipHandler.isEnteringPip(change, info.getType())) {
+                if (pipChange != null) {
+                    throw new IllegalStateException("More than 1 pip-entering changes in one"
+                            + " transition? " + info);
+                }
+                pipChange = change;
+                // going backwards, so remove-by-index is fine.
+                everythingElse.getChanges().remove(i);
+            } else if (isHomeOpening(change)) {
+                homeIsOpening = true;
+            } else if (isWallpaper(change)) {
+                wallpaper = change;
+            }
+        }
+        if (pipChange == null) {
+            // um, something probably went wrong.
+            return false;
+        }
+        final boolean isGoingHome = homeIsOpening;
+        Transitions.TransitionFinishCallback finishCB = (wct) -> {
+            --mixed.mInFlightSubAnimations;
+            mixed.joinFinishArgs(wct);
+            if (mixed.mInFlightSubAnimations > 0) return;
+            if (isGoingHome) {
+                splitHandler.onTransitionAnimationComplete();
+            }
+            finishCallback.onTransitionFinished(mixed.mFinishWCT);
+        };
+        if (isGoingHome || splitHandler.getSplitItemPosition(pipChange.getLastParent())
+                != SPLIT_POSITION_UNDEFINED) {
+            ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, " Animation is actually mixed "
+                    + "since entering-PiP caused us to leave split and return home.");
+            // We need to split the transition into 2 parts: the pip part (animated by pip)
+            // and the dismiss-part (animated by launcher).
+            mixed.mInFlightSubAnimations = 2;
+            // immediately make the wallpaper visible (so that we don't see it pop-in during
+            // the time it takes to start recents animation (which is remote).
+            if (wallpaper != null) {
+                startTransaction.show(wallpaper.getLeash()).setAlpha(wallpaper.getLeash(), 1.f);
+            }
+            // make a new startTransaction because pip's startEnterAnimation "consumes" it so
+            // we need a separate one to send over to launcher.
+            SurfaceControl.Transaction otherStartT = new SurfaceControl.Transaction();
+            @SplitScreen.StageType int topStageToKeep = STAGE_TYPE_UNDEFINED;
+            if (splitHandler.isSplitScreenVisible()) {
+                // The non-going home case, we could be pip-ing one of the split stages and keep
+                // showing the other
+                for (int i = info.getChanges().size() - 1; i >= 0; --i) {
+                    TransitionInfo.Change change = info.getChanges().get(i);
+                    if (change == pipChange) {
+                        // Ignore the change/task that's going into Pip
+                        continue;
+                    }
+                    @SplitScreen.StageType int splitItemStage =
+                            splitHandler.getSplitItemStage(change.getLastParent());
+                    if (splitItemStage != STAGE_TYPE_UNDEFINED) {
+                        topStageToKeep = splitItemStage;
+                        break;
+                    }
+                }
+            }
+            // Let split update internal state for dismiss.
+            splitHandler.prepareDismissAnimation(topStageToKeep,
+                    EXIT_REASON_CHILD_TASK_ENTER_PIP, everythingElse, otherStartT,
+                    finishTransaction);
+
+            // We are trying to accommodate launcher's close animation which can't handle the
+            // divider-bar, so if split-handler is closing the divider-bar, just hide it and
+            // remove from transition info.
+            for (int i = everythingElse.getChanges().size() - 1; i >= 0; --i) {
+                if ((everythingElse.getChanges().get(i).getFlags() & FLAG_IS_DIVIDER_BAR)
+                        != 0) {
+                    everythingElse.getChanges().remove(i);
+                    break;
+                }
+            }
+
+            pipHandler.setEnterAnimationType(ANIM_TYPE_ALPHA);
+            pipHandler.startEnterAnimation(pipChange, startTransaction, finishTransaction,
+                    finishCB);
+            // Dispatch the rest of the transition normally. This will most-likely be taken by
+            // recents or default handler.
+            mixed.mLeftoversHandler = player.dispatchTransition(mixed.mTransition, everythingElse,
+                    otherStartT, finishTransaction, finishCB, mixedHandler);
+        } else {
+            ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, "  Not leaving split, so just "
+                    + "forward animation to Pip-Handler.");
+            // This happens if the pip-ing activity is in a multi-activity task (and thus a
+            // new pip task is spawned). In this case, we don't actually exit split so we can
+            // just let pip transition handle the animation verbatim.
+            mixed.mInFlightSubAnimations = 1;
+            pipHandler.startAnimation(
+                    mixed.mTransition, info, startTransaction, finishTransaction, finishCB);
+        }
+        return true;
+    }
+
+    private static boolean isHomeOpening(@NonNull TransitionInfo.Change change) {
+        return change.getTaskInfo() != null
+                && change.getTaskInfo().getActivityType() == ACTIVITY_TYPE_HOME;
+    }
+
+    private static boolean isWallpaper(@NonNull TransitionInfo.Change change) {
+        return (change.getFlags() & FLAG_IS_WALLPAPER) != 0;
+    }
+
+    static boolean animateKeyguard(
+            @NonNull DefaultMixedHandler.MixedTransition mixed, @NonNull TransitionInfo info,
+            @NonNull SurfaceControl.Transaction startTransaction,
+            @NonNull SurfaceControl.Transaction finishTransaction,
+            @NonNull Transitions.TransitionFinishCallback finishCallback,
+            @NonNull KeyguardTransitionHandler keyguardHandler,
+            PipTransitionController pipHandler) {
+        if (mixed.mFinishT == null) {
+            mixed.mFinishT = finishTransaction;
+            mixed.mFinishCB = finishCallback;
+        }
+        // Sync pip state.
+        if (pipHandler != null) {
+            pipHandler.syncPipSurfaceState(info, startTransaction, finishTransaction);
+        }
+        return mixed.startSubAnimation(keyguardHandler, info, startTransaction, finishTransaction);
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/RecentsMixedTransition.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/RecentsMixedTransition.java
new file mode 100644
index 0000000..643e026
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/RecentsMixedTransition.java
@@ -0,0 +1,214 @@
+/*
+ * 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.transition;
+
+import static android.view.WindowManager.TRANSIT_FLAG_KEYGUARD_UNOCCLUDING;
+
+import static com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_UNDEFINED;
+import static com.android.wm.shell.transition.DefaultMixedHandler.handoverTransitionLeashes;
+import static com.android.wm.shell.transition.MixedTransitionHelper.animateEnterPipFromSplit;
+import static com.android.wm.shell.transition.MixedTransitionHelper.animateKeyguard;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.os.IBinder;
+import android.view.SurfaceControl;
+import android.window.TransitionInfo;
+import android.window.WindowContainerTransaction;
+
+import com.android.wm.shell.desktopmode.DesktopTasksController;
+import com.android.wm.shell.keyguard.KeyguardTransitionHandler;
+import com.android.wm.shell.pip.PipTransitionController;
+import com.android.wm.shell.recents.RecentsTransitionHandler;
+import com.android.wm.shell.splitscreen.StageCoordinator;
+
+class RecentsMixedTransition extends DefaultMixedHandler.MixedTransition {
+    private final RecentsTransitionHandler mRecentsHandler;
+    private final DesktopTasksController mDesktopTasksController;
+
+    RecentsMixedTransition(int type, IBinder transition, Transitions player,
+            DefaultMixedHandler mixedHandler, PipTransitionController pipHandler,
+            StageCoordinator splitHandler, KeyguardTransitionHandler keyguardHandler,
+            RecentsTransitionHandler recentsHandler,
+            DesktopTasksController desktopTasksController) {
+        super(type, transition, player, mixedHandler, pipHandler, splitHandler, keyguardHandler);
+        mRecentsHandler = recentsHandler;
+        mDesktopTasksController = desktopTasksController;
+        mLeftoversHandler = mRecentsHandler;
+    }
+
+    @Override
+    boolean startAnimation(
+            @NonNull IBinder transition, @NonNull TransitionInfo info,
+            @NonNull SurfaceControl.Transaction startTransaction,
+            @NonNull SurfaceControl.Transaction finishTransaction,
+            @NonNull Transitions.TransitionFinishCallback finishCallback) {
+        return switch (mType) {
+            case TYPE_RECENTS_DURING_DESKTOP ->
+                    animateRecentsDuringDesktop(
+                            info, startTransaction, finishTransaction, finishCallback);
+            case TYPE_RECENTS_DURING_KEYGUARD ->
+                    animateRecentsDuringKeyguard(
+                            info, startTransaction, finishTransaction, finishCallback);
+            case TYPE_RECENTS_DURING_SPLIT ->
+                    animateRecentsDuringSplit(
+                            info, startTransaction, finishTransaction, finishCallback);
+            default -> throw new IllegalStateException(
+                    "Starting Recents mixed animation with unknown or illegal type: " + mType);
+        };
+    }
+
+    private boolean animateRecentsDuringDesktop(
+            @NonNull TransitionInfo info,
+            @NonNull SurfaceControl.Transaction startTransaction,
+            @NonNull SurfaceControl.Transaction finishTransaction,
+            @NonNull Transitions.TransitionFinishCallback finishCallback) {
+        if (mInfo == null) {
+            mInfo = info;
+            mFinishT = finishTransaction;
+            mFinishCB = finishCallback;
+        }
+        Transitions.TransitionFinishCallback finishCB = wct -> {
+            mInFlightSubAnimations--;
+            if (mInFlightSubAnimations == 0) {
+                finishCallback.onTransitionFinished(wct);
+            }
+        };
+
+        mInFlightSubAnimations++;
+        boolean consumed = mRecentsHandler.startAnimation(
+                mTransition, info, startTransaction, finishTransaction, finishCB);
+        if (!consumed) {
+            mInFlightSubAnimations--;
+            return false;
+        }
+        if (mDesktopTasksController != null) {
+            mDesktopTasksController.syncSurfaceState(info, finishTransaction);
+            return true;
+        }
+
+        return false;
+    }
+
+    private boolean animateRecentsDuringKeyguard(
+            @NonNull TransitionInfo info,
+            @NonNull SurfaceControl.Transaction startTransaction,
+            @NonNull SurfaceControl.Transaction finishTransaction,
+            @NonNull Transitions.TransitionFinishCallback finishCallback) {
+        if (mInfo == null) {
+            mInfo = info;
+            mFinishT = finishTransaction;
+            mFinishCB = finishCallback;
+        }
+        return startSubAnimation(mRecentsHandler, info, startTransaction, finishTransaction);
+    }
+
+    private boolean animateRecentsDuringSplit(
+            @NonNull TransitionInfo info,
+            @NonNull SurfaceControl.Transaction startTransaction,
+            @NonNull SurfaceControl.Transaction finishTransaction,
+            @NonNull Transitions.TransitionFinishCallback finishCallback) {
+        for (int i = info.getChanges().size() - 1; i >= 0; --i) {
+            final TransitionInfo.Change change = info.getChanges().get(i);
+            // Pip auto-entering info might be appended to recent transition like pressing
+            // home-key in 3-button navigation. This offers split handler the opportunity to
+            // handle split to pip animation.
+            if (mPipHandler.isEnteringPip(change, info.getType())
+                    && mSplitHandler.getSplitItemPosition(change.getLastParent())
+                    != SPLIT_POSITION_UNDEFINED) {
+                return animateEnterPipFromSplit(this, info, startTransaction, finishTransaction,
+                        finishCallback, mPlayer, mMixedHandler, mPipHandler, mSplitHandler);
+            }
+        }
+
+        // Split-screen is only interested in the recents transition finishing (and merging), so
+        // just wrap finish and start recents animation directly.
+        Transitions.TransitionFinishCallback finishCB = (wct) -> {
+            mInFlightSubAnimations = 0;
+            // If pair-to-pair switching, the post-recents clean-up isn't needed.
+            wct = wct != null ? wct : new WindowContainerTransaction();
+            if (mAnimType != ANIM_TYPE_PAIR_TO_PAIR) {
+                mSplitHandler.onRecentsInSplitAnimationFinish(wct, finishTransaction);
+            } else {
+                // notify pair-to-pair recents animation finish
+                mSplitHandler.onRecentsPairToPairAnimationFinish(wct);
+            }
+            mSplitHandler.onTransitionAnimationComplete();
+            finishCallback.onTransitionFinished(wct);
+        };
+        mInFlightSubAnimations = 1;
+        mSplitHandler.onRecentsInSplitAnimationStart(info);
+        final boolean handled = mLeftoversHandler.startAnimation(
+                mTransition, info, startTransaction, finishTransaction, finishCB);
+        if (!handled) {
+            mSplitHandler.onRecentsInSplitAnimationCanceled();
+        }
+        return handled;
+    }
+
+    @Override
+    void mergeAnimation(
+            @NonNull IBinder transition, @NonNull TransitionInfo info,
+            @NonNull SurfaceControl.Transaction t, @NonNull IBinder mergeTarget,
+            @NonNull Transitions.TransitionFinishCallback finishCallback) {
+        switch (mType) {
+            case TYPE_RECENTS_DURING_DESKTOP:
+                mLeftoversHandler.mergeAnimation(transition, info, t, mergeTarget, finishCallback);
+                return;
+            case TYPE_RECENTS_DURING_KEYGUARD:
+                if ((info.getFlags() & TRANSIT_FLAG_KEYGUARD_UNOCCLUDING) != 0) {
+                    handoverTransitionLeashes(mInfo, info, t, mFinishT);
+                    if (animateKeyguard(
+                            this, info, t, mFinishT, mFinishCB, mKeyguardHandler, mPipHandler)) {
+                        finishCallback.onTransitionFinished(null);
+                    }
+                }
+                mLeftoversHandler.mergeAnimation(transition, info, t, mergeTarget,
+                        finishCallback);
+                return;
+            case TYPE_RECENTS_DURING_SPLIT:
+                if (mSplitHandler.isPendingEnter(transition)) {
+                    // Recents -> enter-split means that we are switching from one pair to
+                    // another pair.
+                    mAnimType = DefaultMixedHandler.MixedTransition.ANIM_TYPE_PAIR_TO_PAIR;
+                }
+                mLeftoversHandler.mergeAnimation(transition, info, t, mergeTarget, finishCallback);
+                return;
+            default:
+                throw new IllegalStateException("Playing a Recents mixed transition with unknown or"
+                        + " illegal type: " + mType);
+        }
+    }
+
+    @Override
+    void onTransitionConsumed(
+            @NonNull IBinder transition, boolean aborted,
+            @Nullable SurfaceControl.Transaction finishT) {
+        switch (mType) {
+            case TYPE_RECENTS_DURING_DESKTOP:
+            case TYPE_RECENTS_DURING_SPLIT:
+                mLeftoversHandler.onTransitionConsumed(transition, aborted, finishT);
+                break;
+            default:
+                break;
+        }
+
+        if (mHasRequestToRemote) {
+            mPlayer.getRemoteTransitionHandler().onTransitionConsumed(transition, aborted, finishT);
+        }
+    }
+}
diff --git a/packages/CredentialManager/src/com/android/credentialmanager/autofill/CredentialAutofillService.kt b/packages/CredentialManager/src/com/android/credentialmanager/autofill/CredentialAutofillService.kt
index 58467af..03ac605 100644
--- a/packages/CredentialManager/src/com/android/credentialmanager/autofill/CredentialAutofillService.kt
+++ b/packages/CredentialManager/src/com/android/credentialmanager/autofill/CredentialAutofillService.kt
@@ -173,7 +173,7 @@
                 CancellationSignal(),
                 Executors.newSingleThreadExecutor(),
                 outcome,
-                autofillCallback
+                autofillCallback.asBinder()
         )
     }
 
@@ -358,8 +358,8 @@
                 } else {
                     spec = inlinePresentationSpecs[inlinePresentationSpecsCount - 1]
                 }
-                val displayName: String = if (primaryEntry.credentialType == CredentialType.PASSKEY
-                        && primaryEntry.displayName != null) {
+                val displayName: String = if (primaryEntry.credentialType ==
+                        CredentialType.PASSKEY && primaryEntry.displayName != null) {
                     primaryEntry.displayName!!
                 } else {
                     primaryEntry.userName
diff --git a/packages/SettingsLib/SelectorWithWidgetPreference/res/layout-v33/preference_selector_with_widget.xml b/packages/SettingsLib/SelectorWithWidgetPreference/res/layout-v33/preference_selector_with_widget.xml
index 0b27464..0764609 100644
--- a/packages/SettingsLib/SelectorWithWidgetPreference/res/layout-v33/preference_selector_with_widget.xml
+++ b/packages/SettingsLib/SelectorWithWidgetPreference/res/layout-v33/preference_selector_with_widget.xml
@@ -66,6 +66,7 @@
             android:layout_width="wrap_content"
             android:layout_height="wrap_content"
             android:maxLines="2"
+            android:ellipsize="end"
             android:hyphenationFrequency="normalFast"
             android:lineBreakWordStyle="phrase"
             android:textAppearance="?android:attr/textAppearanceListItem"/>
diff --git a/packages/SettingsLib/SelectorWithWidgetPreference/res/layout/preference_selector_with_widget.xml b/packages/SettingsLib/SelectorWithWidgetPreference/res/layout/preference_selector_with_widget.xml
index 8bb56ff..4f1a910 100644
--- a/packages/SettingsLib/SelectorWithWidgetPreference/res/layout/preference_selector_with_widget.xml
+++ b/packages/SettingsLib/SelectorWithWidgetPreference/res/layout/preference_selector_with_widget.xml
@@ -66,6 +66,7 @@
             android:layout_width="wrap_content"
             android:layout_height="wrap_content"
             android:maxLines="2"
+            android:ellipsize="end"
             android:textAppearance="?android:attr/textAppearanceListItem"/>
 
         <LinearLayout
diff --git a/packages/SettingsLib/Spa/gallery/AndroidManifest.xml b/packages/SettingsLib/Spa/gallery/AndroidManifest.xml
index 965fdcf..ed90284 100644
--- a/packages/SettingsLib/Spa/gallery/AndroidManifest.xml
+++ b/packages/SettingsLib/Spa/gallery/AndroidManifest.xml
@@ -73,6 +73,11 @@
             android:authorities="com.android.spa.gallery.debug.provider"
             android:exported="false">
         </provider>
-
+        <activity
+            android:name="com.android.settingslib.spa.gallery.SpaDialogActivity"
+            android:excludeFromRecents="true"
+            android:exported="true"
+            android:theme="@style/Theme.SpaLib.Dialog">
+        </activity>
     </application>
 </manifest>
diff --git a/packages/SettingsLib/Spa/gallery/src/com/android/settingslib/spa/gallery/SpaDialogActivity.kt b/packages/SettingsLib/Spa/gallery/src/com/android/settingslib/spa/gallery/SpaDialogActivity.kt
new file mode 100644
index 0000000..8b80fe2
--- /dev/null
+++ b/packages/SettingsLib/Spa/gallery/src/com/android/settingslib/spa/gallery/SpaDialogActivity.kt
@@ -0,0 +1,131 @@
+/*
+ * 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.settingslib.spa.gallery
+
+import android.os.Bundle
+import androidx.activity.ComponentActivity
+import androidx.activity.compose.setContent
+import androidx.compose.foundation.layout.width
+import androidx.compose.material.icons.Icons
+import androidx.compose.material.icons.filled.WarningAmber
+import androidx.compose.material3.AlertDialog
+import androidx.compose.material3.Button
+import androidx.compose.material3.Icon
+import androidx.compose.material3.OutlinedButton
+import androidx.compose.material3.Text
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.LaunchedEffect
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.remember
+import androidx.compose.runtime.setValue
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.graphics.vector.ImageVector
+import com.android.settingslib.spa.framework.common.LogCategory
+import com.android.settingslib.spa.framework.common.SpaEnvironmentFactory
+import com.android.settingslib.spa.framework.theme.SettingsTheme
+import com.android.settingslib.spa.widget.dialog.getDialogWidth
+
+
+class SpaDialogActivity : ComponentActivity() {
+    private val spaEnvironment get() = SpaEnvironmentFactory.instance
+
+    override fun onCreate(savedInstanceState: Bundle?) {
+        super.onCreate(savedInstanceState)
+        spaEnvironment.logger.message(TAG, "onCreate", category = LogCategory.FRAMEWORK)
+        setContent {
+            SettingsTheme {
+                Content()
+            }
+        }
+    }
+
+    @Composable
+    fun Content() {
+        var openAlertDialog by remember { mutableStateOf(false) }
+        AlertDialog(openAlertDialog)
+        LaunchedEffect(key1 = Unit) {
+            openAlertDialog = true
+        }
+    }
+
+    @Composable
+    fun AlertDialog(openAlertDialog: Boolean) {
+        when {
+            openAlertDialog -> {
+                AlertDialogExample(
+                    onDismissRequest = { finish() },
+                    onConfirmation = { finish() },
+                    dialogTitle = intent.getStringExtra(DIALOG_TITLE) ?: DIALOG_TITLE,
+                    dialogText = intent.getStringExtra(DIALOG_TEXT) ?: DIALOG_TEXT,
+                    icon = Icons.Default.WarningAmber
+                )
+            }
+        }
+    }
+
+    @Composable
+    fun AlertDialogExample(
+        onDismissRequest: () -> Unit,
+        onConfirmation: () -> Unit,
+        dialogTitle: String,
+        dialogText: String,
+        icon: ImageVector,
+    ) {
+        AlertDialog(
+            modifier = Modifier.width(getDialogWidth()),
+            icon = {
+                Icon(icon, contentDescription = null)
+            },
+            title = {
+                Text(text = dialogTitle)
+            },
+            text = {
+                Text(text = dialogText)
+            },
+            onDismissRequest = {
+                onDismissRequest()
+            },
+            dismissButton = {
+                OutlinedButton(
+                    onClick = {
+                        onDismissRequest()
+                    }
+                ) {
+                    Text(intent.getStringExtra(DISMISS_TEXT) ?: DISMISS_TEXT)
+                }
+            },
+            confirmButton = {
+                Button(
+                    onClick = {
+                        onConfirmation()
+                    },
+                ) {
+                    Text(intent.getStringExtra(CONFIRM_TEXT) ?: CONFIRM_TEXT)
+                }
+            }
+        )
+    }
+
+    companion object {
+        private const val TAG = "SpaDialogActivity"
+        private const val DIALOG_TITLE = "dialogTitle"
+        private const val DIALOG_TEXT = "dialogText"
+        private const val CONFIRM_TEXT = "confirmText"
+        private const val DISMISS_TEXT = "dismissText"
+    }
+}
diff --git a/packages/SettingsLib/Spa/spa/res/values/themes.xml b/packages/SettingsLib/Spa/spa/res/values/themes.xml
index 25846ec..4b5a9bc 100644
--- a/packages/SettingsLib/Spa/spa/res/values/themes.xml
+++ b/packages/SettingsLib/Spa/spa/res/values/themes.xml
@@ -22,4 +22,6 @@
         <item name="android:windowActionBar">false</item>
         <item name="android:windowNoTitle">true</item>
     </style>
+
+    <style name="Theme.SpaLib.Dialog" parent="Theme.Material3.DayNight.Dialog"/>
 </resources>
diff --git a/packages/SettingsLib/Spa/spa/src/com/android/settingslib/spa/widget/dialog/SettingsAlertDialog.kt b/packages/SettingsLib/Spa/spa/src/com/android/settingslib/spa/widget/dialog/SettingsAlertDialog.kt
index 207c174..8ffd799 100644
--- a/packages/SettingsLib/Spa/spa/src/com/android/settingslib/spa/widget/dialog/SettingsAlertDialog.kt
+++ b/packages/SettingsLib/Spa/spa/src/com/android/settingslib/spa/widget/dialog/SettingsAlertDialog.kt
@@ -99,7 +99,7 @@
 }
 
 @Composable
-private fun getDialogWidth(): Dp {
+fun getDialogWidth(): Dp {
     val configuration = LocalConfiguration.current
     return configuration.screenWidthDp.dp * when (configuration.orientation) {
         Configuration.ORIENTATION_LANDSCAPE -> 0.6f
diff --git a/packages/SettingsLib/src/com/android/settingslib/Utils.java b/packages/SettingsLib/src/com/android/settingslib/Utils.java
index c2be571..fb14a17 100644
--- a/packages/SettingsLib/src/com/android/settingslib/Utils.java
+++ b/packages/SettingsLib/src/com/android/settingslib/Utils.java
@@ -1,6 +1,7 @@
 package com.android.settingslib;
 
 import static android.app.admin.DevicePolicyResources.Strings.Settings.WORK_PROFILE_USER_LABEL;
+import static android.webkit.Flags.updateServiceV2;
 
 import android.annotation.ColorInt;
 import android.app.admin.DevicePolicyManager;
@@ -34,6 +35,7 @@
 import android.net.wifi.WifiInfo;
 import android.os.BatteryManager;
 import android.os.Build;
+import android.os.RemoteException;
 import android.os.SystemProperties;
 import android.os.UserHandle;
 import android.os.UserManager;
@@ -44,6 +46,9 @@
 import android.telephony.ServiceState;
 import android.telephony.TelephonyManager;
 import android.util.Log;
+import android.webkit.IWebViewUpdateService;
+import android.webkit.WebViewFactory;
+import android.webkit.WebViewProviderInfo;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -65,6 +70,8 @@
 
 public class Utils {
 
+    private static final String TAG = "Utils";
+
     @VisibleForTesting
     static final String STORAGE_MANAGER_ENABLED_PROPERTY =
             "ro.storage_manager.enabled";
@@ -76,6 +83,7 @@
     private static String sPermissionControllerPackageName;
     private static String sServicesSystemSharedLibPackageName;
     private static String sSharedSystemSharedLibPackageName;
+    private static String sDefaultWebViewPackageName;
 
     static final int[] WIFI_PIE = {
         com.android.internal.R.drawable.ic_wifi_signal_0,
@@ -445,6 +453,7 @@
                 || packageName.equals(sServicesSystemSharedLibPackageName)
                 || packageName.equals(sSharedSystemSharedLibPackageName)
                 || packageName.equals(PrintManager.PRINT_SPOOLER_PACKAGE_NAME)
+                || (updateServiceV2() && packageName.equals(getDefaultWebViewPackageName()))
                 || isDeviceProvisioningPackage(resources, packageName);
     }
 
@@ -459,6 +468,29 @@
     }
 
     /**
+     * Fetch the package name of the default WebView provider.
+     */
+    @Nullable
+    private static String getDefaultWebViewPackageName() {
+        if (sDefaultWebViewPackageName != null) {
+            return sDefaultWebViewPackageName;
+        }
+
+        try {
+            IWebViewUpdateService service = WebViewFactory.getUpdateService();
+            if (service != null) {
+                WebViewProviderInfo provider = service.getDefaultWebViewPackage();
+                if (provider != null) {
+                    sDefaultWebViewPackageName = provider.packageName;
+                }
+            }
+        } catch (RemoteException e) {
+            Log.e(TAG, "RemoteException when trying to fetch default WebView package Name", e);
+        }
+        return sDefaultWebViewPackageName;
+    }
+
+    /**
      * Returns the Wifi icon resource for a given RSSI level.
      *
      * @param level The number of bars to show (0-4)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/keyguard/KeyguardSecurityContainerControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/keyguard/KeyguardSecurityContainerControllerTest.kt
index 02d30c5e..bc3ca1b 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/keyguard/KeyguardSecurityContainerControllerTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/keyguard/KeyguardSecurityContainerControllerTest.kt
@@ -230,11 +230,7 @@
         sceneTransitionStateFlow =
             MutableStateFlow(ObservableTransitionState.Idle(SceneKey.Lockscreen))
         sceneInteractor.setTransitionState(sceneTransitionStateFlow)
-        deviceEntryInteractor =
-            sceneTestUtils.deviceEntryInteractor(
-                authenticationInteractor = sceneTestUtils.authenticationInteractor(),
-                sceneInteractor = sceneInteractor,
-            )
+        deviceEntryInteractor = sceneTestUtils.deviceEntryInteractor()
 
         mSetFlagsRule.disableFlags(FLAG_SIDEFPS_CONTROLLER_REFACTOR)
         underTest =
@@ -253,7 +249,7 @@
                 falsingManager,
                 userSwitcherController,
                 featureFlags,
-                sceneTestUtils.sceneContainerFlags,
+                sceneTestUtils.fakeSceneContainerFlags,
                 globalSettings,
                 sessionTracker,
                 Optional.of(sideFpsController),
@@ -791,6 +787,7 @@
     @Test
     fun dismissesKeyguard_whenSceneChangesToGone() =
         sceneTestUtils.testScope.runTest {
+            sceneTestUtils.fakeSceneContainerFlags.enabled = true
             // Upon init, we have never dismisses the keyguard.
             underTest.onInit()
             runCurrent()
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/authentication/data/repository/AuthenticationRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/authentication/data/repository/AuthenticationRepositoryTest.kt
index b9759cc..b4d4e1f 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/authentication/data/repository/AuthenticationRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/authentication/data/repository/AuthenticationRepositoryTest.kt
@@ -83,7 +83,7 @@
             AuthenticationRepositoryImpl(
                 applicationScope = testScope.backgroundScope,
                 backgroundDispatcher = testUtils.testDispatcher,
-                flags = testUtils.sceneContainerFlags,
+                flags = testUtils.fakeSceneContainerFlags,
                 clock = clock,
                 getSecurityMode = getSecurityMode,
                 userRepository = userRepository,
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/domain/interactor/BouncerActionButtonInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/domain/interactor/BouncerActionButtonInteractorTest.kt
index 67ce86b..63581b3 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/domain/interactor/BouncerActionButtonInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/domain/interactor/BouncerActionButtonInteractorTest.kt
@@ -16,27 +16,25 @@
 
 package com.android.systemui.bouncer.domain.interactor
 
-import android.app.ActivityTaskManager
 import android.telecom.TelecomManager
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.app.activityTaskManager
 import com.android.internal.R
+import com.android.internal.logging.fakeMetricsLogger
 import com.android.internal.logging.nano.MetricsProto
-import com.android.internal.logging.testing.FakeMetricsLogger
-import com.android.internal.util.EmergencyAffordanceManager
+import com.android.internal.util.emergencyAffordanceManager
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.flags.Flags.REFACTOR_GETCURRENTUSER
-import com.android.systemui.log.table.TableLogBuffer
 import com.android.systemui.scene.SceneTestUtils
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionsRepository
-import com.android.systemui.statusbar.pipeline.mobile.util.FakeMobileMappingsProxy
 import com.android.systemui.user.domain.interactor.SelectedUserInteractor
 import com.android.systemui.util.mockito.whenever
+import com.android.telecom.telecomManager
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
@@ -53,27 +51,26 @@
 @RunWith(AndroidJUnit4::class)
 class BouncerActionButtonInteractorTest : SysuiTestCase() {
 
-    @Mock private lateinit var activityTaskManager: ActivityTaskManager
-    @Mock private lateinit var emergencyAffordanceManager: EmergencyAffordanceManager
     @Mock private lateinit var selectedUserInteractor: SelectedUserInteractor
-    @Mock private lateinit var tableLogger: TableLogBuffer
     @Mock private lateinit var telecomManager: TelecomManager
 
-    private lateinit var utils: SceneTestUtils
-    private lateinit var testScope: TestScope
+    private val utils = SceneTestUtils(this)
+    private val testScope = utils.testScope
+    private val metricsLogger = utils.kosmos.fakeMetricsLogger
+    private val activityTaskManager = utils.kosmos.activityTaskManager
+    private val emergencyAffordanceManager = utils.kosmos.emergencyAffordanceManager
+
     private lateinit var mobileConnectionsRepository: FakeMobileConnectionsRepository
 
-    private val metricsLogger = FakeMetricsLogger()
     private var currentUserId: Int = 0
     private var needsEmergencyAffordance = true
 
-    private lateinit var underTest: BouncerActionButtonInteractor
-
     @Before
     fun setUp() {
-        utils = SceneTestUtils(this)
-        testScope = utils.testScope
         MockitoAnnotations.initMocks(this)
+        utils.fakeSceneContainerFlags.enabled = true
+
+        mobileConnectionsRepository = utils.mobileConnectionsRepository
 
         overrideResource(R.string.lockscreen_emergency_call, MESSAGE_EMERGENCY_CALL)
         overrideResource(R.string.lockscreen_return_to_call, MESSAGE_RETURN_TO_CALL)
@@ -86,34 +83,18 @@
             .thenReturn(needsEmergencyAffordance)
         whenever(telecomManager.isInCall).thenReturn(false)
 
-        utils.featureFlags.set(REFACTOR_GETCURRENTUSER, true)
-
-        mobileConnectionsRepository =
-            FakeMobileConnectionsRepository(FakeMobileMappingsProxy(), tableLogger)
+        utils.fakeFeatureFlags.set(REFACTOR_GETCURRENTUSER, true)
 
         utils.telephonyRepository.setHasTelephonyRadio(true)
 
-        underTest =
-            utils.bouncerActionButtonInteractor(
-                mobileConnectionsRepository = mobileConnectionsRepository,
-                activityTaskManager = activityTaskManager,
-                telecomManager = telecomManager,
-                emergencyAffordanceManager = emergencyAffordanceManager,
-                metricsLogger = metricsLogger,
-            )
+        utils.kosmos.telecomManager = telecomManager
     }
 
     @Test
     fun noTelephonyRadio_noButton() =
         testScope.runTest {
             utils.telephonyRepository.setHasTelephonyRadio(false)
-            underTest =
-                utils.bouncerActionButtonInteractor(
-                    mobileConnectionsRepository = mobileConnectionsRepository,
-                    activityTaskManager = activityTaskManager,
-                    telecomManager = telecomManager,
-                )
-
+            val underTest = utils.bouncerActionButtonInteractor()
             val actionButton by collectLastValue(underTest.actionButton)
             assertThat(actionButton).isNull()
         }
@@ -121,12 +102,8 @@
     @Test
     fun noTelecomManager_noButton() =
         testScope.runTest {
-            underTest =
-                utils.bouncerActionButtonInteractor(
-                    mobileConnectionsRepository = mobileConnectionsRepository,
-                    activityTaskManager = activityTaskManager,
-                    telecomManager = null,
-                )
+            utils.kosmos.telecomManager = null
+            val underTest = utils.bouncerActionButtonInteractor()
             val actionButton by collectLastValue(underTest.actionButton)
             assertThat(actionButton).isNull()
         }
@@ -134,6 +111,7 @@
     @Test
     fun duringCall_returnToCallButton() =
         testScope.runTest {
+            val underTest = utils.bouncerActionButtonInteractor()
             val actionButton by collectLastValue(underTest.actionButton)
             utils.telephonyRepository.setIsInCall(true)
 
@@ -143,6 +121,7 @@
             assertThat(actionButton?.onLongClick).isNull()
 
             actionButton?.onClick?.invoke()
+            runCurrent()
 
             assertThat(metricsLogger.logs.size).isEqualTo(1)
             assertThat(metricsLogger.logs.element().category)
@@ -154,6 +133,7 @@
     @Test
     fun noCall_secureAuthMethod_emergencyCallButton() =
         testScope.runTest {
+            val underTest = utils.bouncerActionButtonInteractor()
             val actionButton by collectLastValue(underTest.actionButton)
             mobileConnectionsRepository.isAnySimSecure.value = false
             utils.authenticationRepository.setAuthenticationMethod(AuthenticationMethodModel.Pin)
@@ -165,6 +145,7 @@
             assertThat(actionButton?.onLongClick).isNotNull()
 
             actionButton?.onClick?.invoke()
+            runCurrent()
 
             assertThat(metricsLogger.logs.size).isEqualTo(1)
             assertThat(metricsLogger.logs.element().category)
@@ -182,10 +163,12 @@
     @Test
     fun noCall_insecureAuthMethodButSecureSim_emergencyCallButton() =
         testScope.runTest {
+            val underTest = utils.bouncerActionButtonInteractor()
             val actionButton by collectLastValue(underTest.actionButton)
             mobileConnectionsRepository.isAnySimSecure.value = true
             utils.authenticationRepository.setAuthenticationMethod(AuthenticationMethodModel.None)
             utils.telephonyRepository.setIsInCall(false)
+            runCurrent()
 
             assertThat(actionButton).isNotNull()
             assertThat(actionButton?.label).isEqualTo(MESSAGE_EMERGENCY_CALL)
@@ -196,6 +179,7 @@
     @Test
     fun noCall_insecure_noButton() =
         testScope.runTest {
+            val underTest = utils.bouncerActionButtonInteractor()
             val actionButton by collectLastValue(underTest.actionButton)
             mobileConnectionsRepository.isAnySimSecure.value = false
             utils.authenticationRepository.setAuthenticationMethod(AuthenticationMethodModel.None)
@@ -207,6 +191,7 @@
     @Test
     fun noCall_simSecureButEmergencyNotSupported_noButton() =
         testScope.runTest {
+            val underTest = utils.bouncerActionButtonInteractor()
             val actionButton by collectLastValue(underTest.actionButton)
             mobileConnectionsRepository.isAnySimSecure.value = true
             overrideResource(R.bool.config_enable_emergency_call_while_sim_locked, false)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/domain/interactor/BouncerInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/domain/interactor/BouncerInteractorTest.kt
index 93ba6a4..4b6199b 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/domain/interactor/BouncerInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/domain/interactor/BouncerInteractorTest.kt
@@ -25,7 +25,8 @@
 import com.android.systemui.authentication.shared.model.AuthenticationPatternCoordinate
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.coroutines.collectValues
-import com.android.systemui.deviceentry.domain.interactor.DeviceEntryFaceAuthInteractor
+import com.android.systemui.deviceentry.domain.interactor.deviceEntryFaceAuthInteractor
+import com.android.systemui.keyguard.data.repository.fakeDeviceEntryFaceAuthRepository
 import com.android.systemui.res.R
 import com.android.systemui.scene.SceneTestUtils
 import com.google.common.truth.Truth.assertThat
@@ -37,8 +38,6 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mock
-import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
 
 @OptIn(ExperimentalCoroutinesApi::class)
@@ -46,9 +45,7 @@
 @RunWith(AndroidJUnit4::class)
 class BouncerInteractorTest : SysuiTestCase() {
 
-    @Mock private lateinit var mDeviceEntryFaceAuthInteractor: DeviceEntryFaceAuthInteractor
-
-    private val utils = SceneTestUtils(this)
+    private val utils = SceneTestUtils(this).apply { fakeSceneContainerFlags.enabled = true }
     private val testScope = utils.testScope
     private val authenticationInteractor = utils.authenticationInteractor()
 
@@ -57,6 +54,7 @@
     @Before
     fun setUp() {
         MockitoAnnotations.initMocks(this)
+
         overrideResource(R.string.keyguard_enter_your_pin, MESSAGE_ENTER_YOUR_PIN)
         overrideResource(R.string.keyguard_enter_your_password, MESSAGE_ENTER_YOUR_PASSWORD)
         overrideResource(R.string.keyguard_enter_your_pattern, MESSAGE_ENTER_YOUR_PATTERN)
@@ -64,11 +62,7 @@
         overrideResource(R.string.kg_wrong_password, MESSAGE_WRONG_PASSWORD)
         overrideResource(R.string.kg_wrong_pattern, MESSAGE_WRONG_PATTERN)
 
-        underTest =
-            utils.bouncerInteractor(
-                authenticationInteractor = authenticationInteractor,
-                deviceEntryFaceAuthInteractor = mDeviceEntryFaceAuthInteractor,
-            )
+        underTest = utils.bouncerInteractor()
     }
 
     @Test
@@ -305,8 +299,16 @@
     @Test
     fun intentionalUserInputEvent_notifiesFaceAuthInteractor() =
         testScope.runTest {
+            val isFaceAuthRunning by
+                collectLastValue(utils.kosmos.fakeDeviceEntryFaceAuthRepository.isAuthRunning)
+            utils.kosmos.deviceEntryFaceAuthInteractor.onDeviceLifted()
+            runCurrent()
+            assertThat(isFaceAuthRunning).isTrue()
+
             underTest.onIntentionalUserInput()
-            verify(mDeviceEntryFaceAuthInteractor).onPrimaryBouncerUserInput()
+            runCurrent()
+
+            assertThat(isFaceAuthRunning).isFalse()
         }
 
     companion object {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/AuthMethodBouncerViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/AuthMethodBouncerViewModelTest.kt
index 2f0843b..3043a71 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/AuthMethodBouncerViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/AuthMethodBouncerViewModelTest.kt
@@ -35,10 +35,7 @@
 
     private val utils = SceneTestUtils(this)
     private val testScope = utils.testScope
-    private val bouncerInteractor =
-        utils.bouncerInteractor(
-            authenticationInteractor = utils.authenticationInteractor(),
-        )
+    private val bouncerInteractor = utils.bouncerInteractor()
     private val underTest =
         PinBouncerViewModel(
             applicationContext = context,
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModelTest.kt
index 47bbe6f4..4b1f9fe 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModelTest.kt
@@ -41,6 +41,7 @@
 import kotlinx.coroutines.test.currentTime
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
+import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 
@@ -52,15 +53,14 @@
     private val utils = SceneTestUtils(this)
     private val testScope = utils.testScope
     private val authenticationInteractor = utils.authenticationInteractor()
-    private val bouncerInteractor =
-        utils.bouncerInteractor(
-            authenticationInteractor = authenticationInteractor,
-        )
-    private val underTest =
-        utils.bouncerViewModel(
-            bouncerInteractor = bouncerInteractor,
-            authenticationInteractor = authenticationInteractor,
-        )
+    private val bouncerInteractor = utils.bouncerInteractor()
+    private lateinit var underTest: BouncerViewModel
+
+    @Before
+    fun setUp() {
+        utils.fakeSceneContainerFlags.enabled = true
+        underTest = utils.bouncerViewModel()
+    }
 
     @Test
     fun authMethod_nonNullForSecureMethods_nullForNotSecureMethods() =
@@ -228,13 +228,13 @@
     fun isSideBySideSupported() =
         testScope.runTest {
             val isSideBySideSupported by collectLastValue(underTest.isSideBySideSupported)
-            utils.featureFlags.set(Flags.FULL_SCREEN_USER_SWITCHER, true)
+            utils.fakeFeatureFlags.set(Flags.FULL_SCREEN_USER_SWITCHER, true)
             utils.authenticationRepository.setAuthenticationMethod(Pin)
             assertThat(isSideBySideSupported).isTrue()
             utils.authenticationRepository.setAuthenticationMethod(Password)
             assertThat(isSideBySideSupported).isTrue()
 
-            utils.featureFlags.set(Flags.FULL_SCREEN_USER_SWITCHER, false)
+            utils.fakeFeatureFlags.set(Flags.FULL_SCREEN_USER_SWITCHER, false)
             utils.authenticationRepository.setAuthenticationMethod(Pin)
             assertThat(isSideBySideSupported).isTrue()
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PasswordBouncerViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PasswordBouncerViewModelTest.kt
index 64e6e57..5c5632f 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PasswordBouncerViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PasswordBouncerViewModelTest.kt
@@ -47,16 +47,8 @@
     private val testScope = utils.testScope
     private val authenticationInteractor = utils.authenticationInteractor()
     private val sceneInteractor = utils.sceneInteractor()
-    private val bouncerInteractor =
-        utils.bouncerInteractor(
-            authenticationInteractor = authenticationInteractor,
-        )
-    private val bouncerViewModel =
-        utils.bouncerViewModel(
-            bouncerInteractor = bouncerInteractor,
-            authenticationInteractor = authenticationInteractor,
-            actionButtonInteractor = utils.bouncerActionButtonInteractor(),
-        )
+    private val bouncerInteractor = utils.bouncerInteractor()
+    private val bouncerViewModel = utils.bouncerViewModel()
     private val isInputEnabled = MutableStateFlow(true)
 
     private val underTest =
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PatternBouncerViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PatternBouncerViewModelTest.kt
index 83d1938..9ee344a 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PatternBouncerViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PatternBouncerViewModelTest.kt
@@ -48,16 +48,8 @@
     private val testScope = utils.testScope
     private val authenticationInteractor = utils.authenticationInteractor()
     private val sceneInteractor = utils.sceneInteractor()
-    private val bouncerInteractor =
-        utils.bouncerInteractor(
-            authenticationInteractor = authenticationInteractor,
-        )
-    private val bouncerViewModel =
-        utils.bouncerViewModel(
-            bouncerInteractor = bouncerInteractor,
-            authenticationInteractor = authenticationInteractor,
-            actionButtonInteractor = utils.bouncerActionButtonInteractor(),
-        )
+    private val bouncerInteractor = utils.bouncerInteractor()
+    private val bouncerViewModel = utils.bouncerViewModel()
     private val underTest =
         PatternBouncerViewModel(
             applicationContext = context,
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PinBouncerViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PinBouncerViewModelTest.kt
index db98d76..75e372f 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PinBouncerViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PinBouncerViewModelTest.kt
@@ -47,16 +47,8 @@
     private val testScope = utils.testScope
     private val sceneInteractor = utils.sceneInteractor()
     private val authenticationInteractor = utils.authenticationInteractor()
-    private val bouncerInteractor =
-        utils.bouncerInteractor(
-            authenticationInteractor = authenticationInteractor,
-        )
-    private val bouncerViewModel =
-        utils.bouncerViewModel(
-            bouncerInteractor = bouncerInteractor,
-            authenticationInteractor = authenticationInteractor,
-            actionButtonInteractor = utils.bouncerActionButtonInteractor(),
-        )
+    private val bouncerInteractor = utils.bouncerInteractor()
+    private val bouncerViewModel = utils.bouncerViewModel()
     private val underTest =
         PinBouncerViewModel(
             applicationContext = context,
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractorTest.kt
index ea19cb7..929e879 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractorTest.kt
@@ -23,9 +23,8 @@
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.coroutines.collectValues
-import com.android.systemui.deviceentry.data.repository.FakeDeviceEntryRepository
-import com.android.systemui.keyguard.data.repository.FakeDeviceEntryFaceAuthRepository
-import com.android.systemui.keyguard.data.repository.FakeTrustRepository
+import com.android.systemui.keyguard.data.repository.fakeDeviceEntryFaceAuthRepository
+import com.android.systemui.keyguard.data.repository.fakeTrustRepository
 import com.android.systemui.scene.SceneTestUtils
 import com.android.systemui.scene.shared.model.SceneKey
 import com.android.systemui.scene.shared.model.SceneModel
@@ -33,6 +32,7 @@
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
+import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 
@@ -43,19 +43,17 @@
 
     private val utils = SceneTestUtils(this)
     private val testScope = utils.testScope
-    private val repository: FakeDeviceEntryRepository = utils.deviceEntryRepository
-    private val faceAuthRepository = FakeDeviceEntryFaceAuthRepository()
-    private val trustRepository = FakeTrustRepository()
+    private val faceAuthRepository = utils.kosmos.fakeDeviceEntryFaceAuthRepository
+    private val trustRepository = utils.kosmos.fakeTrustRepository
     private val sceneInteractor = utils.sceneInteractor()
     private val authenticationInteractor = utils.authenticationInteractor()
-    private val underTest =
-        utils.deviceEntryInteractor(
-            repository = repository,
-            authenticationInteractor = authenticationInteractor,
-            sceneInteractor = sceneInteractor,
-            faceAuthRepository = faceAuthRepository,
-            trustRepository = trustRepository,
-        )
+    private lateinit var underTest: DeviceEntryInteractor
+
+    @Before
+    fun setUp() {
+        utils.fakeSceneContainerFlags.enabled = true
+        underTest = utils.deviceEntryInteractor()
+    }
 
     @Test
     fun canSwipeToEnter_startsNull() =
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt
index 11939c1..8933d2c 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt
@@ -63,7 +63,7 @@
             repository = repository,
             commandQueue = commandQueue,
             powerInteractor = PowerInteractorFactory.create().powerInteractor,
-            sceneContainerFlags = testUtils.sceneContainerFlags,
+            sceneContainerFlags = testUtils.fakeSceneContainerFlags,
             bouncerRepository = bouncerRepository,
             configurationInteractor = ConfigurationInteractor(FakeConfigurationRepository()),
             shadeRepository = shadeRepository,
@@ -183,6 +183,7 @@
     @Test
     fun animationDozingTransitions() =
         testScope.runTest {
+            testUtils.fakeSceneContainerFlags.enabled = true
             val isAnimate by collectLastValue(underTest.animateDozingTransitions)
 
             underTest.setAnimateDozingTransitions(true)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenSceneViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenSceneViewModelTest.kt
index 74d309c..04e90c8 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenSceneViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenSceneViewModelTest.kt
@@ -87,11 +87,7 @@
     private fun createLockscreenSceneViewModel(): LockscreenSceneViewModel {
         return LockscreenSceneViewModel(
             applicationScope = testScope.backgroundScope,
-            deviceEntryInteractor =
-                utils.deviceEntryInteractor(
-                    authenticationInteractor = utils.authenticationInteractor(),
-                    sceneInteractor = utils.sceneInteractor(),
-                ),
+            deviceEntryInteractor = utils.deviceEntryInteractor(),
             communalInteractor = utils.communalInteractor(),
             longPress =
                 KeyguardLongPressViewModel(
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt
index 530d127d..ecc2ef1 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt
@@ -24,6 +24,7 @@
 import androidx.test.filters.SmallTest
 import com.android.internal.R
 import com.android.internal.util.EmergencyAffordanceManager
+import com.android.internal.util.emergencyAffordanceManager
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.authentication.data.repository.FakeAuthenticationRepository
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
@@ -35,13 +36,11 @@
 import com.android.systemui.flags.Flags
 import com.android.systemui.keyguard.ui.viewmodel.KeyguardLongPressViewModel
 import com.android.systemui.keyguard.ui.viewmodel.LockscreenSceneViewModel
-import com.android.systemui.log.table.TableLogBuffer
 import com.android.systemui.media.controls.pipeline.MediaDataManager
 import com.android.systemui.media.controls.ui.MediaHost
 import com.android.systemui.model.SysUiState
 import com.android.systemui.power.domain.interactor.PowerInteractor.Companion.setAsleepForTest
 import com.android.systemui.power.domain.interactor.PowerInteractor.Companion.setAwakeForTest
-import com.android.systemui.power.domain.interactor.PowerInteractorFactory
 import com.android.systemui.qs.ui.adapter.FakeQSSceneAdapter
 import com.android.systemui.scene.domain.startable.SceneContainerStartable
 import com.android.systemui.scene.shared.model.ObservableTransitionState
@@ -61,6 +60,7 @@
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
+import com.android.telecom.telecomManager
 import com.google.common.truth.Truth.assertThat
 import com.google.common.truth.Truth.assertWithMessage
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -101,27 +101,12 @@
 @RunWith(AndroidJUnit4::class)
 class SceneFrameworkIntegrationTest : SysuiTestCase() {
 
-    @Mock private lateinit var emergencyAffordanceManager: EmergencyAffordanceManager
-    @Mock private lateinit var tableLogger: TableLogBuffer
-    @Mock private lateinit var telecomManager: TelecomManager
-
-    private val utils = SceneTestUtils(this)
+    private val utils = SceneTestUtils(this).apply { fakeSceneContainerFlags.enabled = true }
     private val testScope = utils.testScope
     private val sceneContainerConfig = utils.fakeSceneContainerConfig()
-    private val sceneRepository =
-        utils.fakeSceneContainerRepository(
-            containerConfig = sceneContainerConfig,
-        )
-    private val sceneInteractor =
-        utils.sceneInteractor(
-            repository = sceneRepository,
-        )
+    private val sceneInteractor = utils.sceneInteractor()
     private val authenticationInteractor = utils.authenticationInteractor()
-    private val deviceEntryInteractor =
-        utils.deviceEntryInteractor(
-            authenticationInteractor = authenticationInteractor,
-            sceneInteractor = sceneInteractor,
-        )
+    private val deviceEntryInteractor = utils.deviceEntryInteractor()
     private val communalInteractor = utils.communalInteractor()
 
     private val transitionState =
@@ -135,10 +120,7 @@
             )
             .apply { setTransitionState(transitionState) }
 
-    private val bouncerInteractor =
-        utils.bouncerInteractor(
-            authenticationInteractor = authenticationInteractor,
-        )
+    private val bouncerInteractor = utils.bouncerInteractor()
 
     private lateinit var mobileConnectionsRepository: FakeMobileConnectionsRepository
     private lateinit var bouncerActionButtonInteractor: BouncerActionButtonInteractor
@@ -170,19 +152,15 @@
                     FakeMobileConnectionsRepository(),
                 ),
             constants = mock(),
-            utils.featureFlags,
+            utils.fakeFeatureFlags,
             scope = testScope.backgroundScope,
         )
 
     private lateinit var shadeHeaderViewModel: ShadeHeaderViewModel
     private lateinit var shadeSceneViewModel: ShadeSceneViewModel
 
-    private val keyguardRepository = utils.keyguardRepository
-    private val keyguardInteractor =
-        utils.keyguardInteractor(
-            repository = keyguardRepository,
-        )
-    private val powerInteractor = PowerInteractorFactory.create().powerInteractor
+    private val keyguardInteractor = utils.keyguardInteractor()
+    private val powerInteractor = utils.powerInteractor()
 
     private var bouncerSceneJob: Job? = null
 
@@ -191,21 +169,26 @@
     @Mock private lateinit var mediaDataManager: MediaDataManager
     @Mock private lateinit var mediaHost: MediaHost
 
+    private lateinit var emergencyAffordanceManager: EmergencyAffordanceManager
+    private lateinit var telecomManager: TelecomManager
+
     @Before
     fun setUp() {
         MockitoAnnotations.initMocks(this)
+
         overrideResource(R.bool.config_enable_emergency_call_while_sim_locked, true)
+        telecomManager = checkNotNull(utils.kosmos.telecomManager)
         whenever(telecomManager.isInCall).thenReturn(false)
+        emergencyAffordanceManager = utils.kosmos.emergencyAffordanceManager
         whenever(emergencyAffordanceManager.needsEmergencyAffordance()).thenReturn(true)
 
-        utils.featureFlags.apply {
+        utils.fakeFeatureFlags.apply {
             set(Flags.NEW_NETWORK_SLICE_UI, false)
             set(Flags.REFACTOR_GETCURRENTUSER, true)
         }
 
-        mobileConnectionsRepository =
-            FakeMobileConnectionsRepository(FakeMobileMappingsProxy(), tableLogger)
-        mobileConnectionsRepository.isAnySimSecure.value = true
+        mobileConnectionsRepository = utils.mobileConnectionsRepository
+        mobileConnectionsRepository.isAnySimSecure.value = false
 
         utils.telephonyRepository.apply {
             setHasTelephonyRadio(true)
@@ -213,18 +196,8 @@
             setIsInCall(false)
         }
 
-        bouncerActionButtonInteractor =
-            utils.bouncerActionButtonInteractor(
-                mobileConnectionsRepository = mobileConnectionsRepository,
-                telecomManager = telecomManager,
-                emergencyAffordanceManager = emergencyAffordanceManager,
-            )
-        bouncerViewModel =
-            utils.bouncerViewModel(
-                bouncerInteractor = bouncerInteractor,
-                authenticationInteractor = authenticationInteractor,
-                actionButtonInteractor = bouncerActionButtonInteractor,
-            )
+        bouncerActionButtonInteractor = utils.bouncerActionButtonInteractor()
+        bouncerViewModel = utils.bouncerViewModel()
 
         shadeHeaderViewModel =
             ShadeHeaderViewModel(
@@ -257,7 +230,7 @@
                 sceneInteractor = sceneInteractor,
                 deviceEntryInteractor = deviceEntryInteractor,
                 keyguardInteractor = keyguardInteractor,
-                flags = utils.sceneContainerFlags,
+                flags = utils.fakeSceneContainerFlags,
                 sysUiState = sysUiState,
                 displayId = displayTracker.defaultDisplayId,
                 sceneLogger = mock(),
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryTest.kt
index ddeb05b..339d026 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryTest.kt
@@ -39,7 +39,7 @@
 @RunWith(AndroidJUnit4::class)
 class SceneContainerRepositoryTest : SysuiTestCase() {
 
-    private val utils = SceneTestUtils(this)
+    private val utils = SceneTestUtils(this).apply { fakeSceneContainerFlags.enabled = true }
     private val testScope = utils.testScope
 
     @Test
@@ -72,7 +72,7 @@
     @Test(expected = IllegalStateException::class)
     fun setDesiredScene_noSuchSceneInContainer_throws() {
         utils.kosmos.sceneKeys = listOf(SceneKey.QuickSettings, SceneKey.Lockscreen)
-        val underTest = utils.fakeSceneContainerRepository(utils.fakeSceneContainerConfig())
+        val underTest = utils.fakeSceneContainerRepository()
         underTest.setDesiredScene(SceneModel(SceneKey.Shade))
     }
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/interactor/SceneInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/interactor/SceneInteractorTest.kt
index 7f4bbbe..486f7ab 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/interactor/SceneInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/interactor/SceneInteractorTest.kt
@@ -28,6 +28,7 @@
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.test.runTest
+import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 
@@ -37,8 +38,14 @@
 
     private val utils = SceneTestUtils(this)
     private val testScope = utils.testScope
-    private val repository = utils.fakeSceneContainerRepository()
-    private val underTest = utils.sceneInteractor(repository = repository)
+
+    private lateinit var underTest: SceneInteractor
+
+    @Before
+    fun setUp() {
+        utils.fakeSceneContainerFlags.enabled = true
+        underTest = utils.sceneInteractor()
+    }
 
     @Test
     fun allSceneKeys() {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/startable/SceneContainerStartableTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/startable/SceneContainerStartableTest.kt
index dd22976..5fe4ca1 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/startable/SceneContainerStartableTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/startable/SceneContainerStartableTest.kt
@@ -28,7 +28,7 @@
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
 import com.android.systemui.classifier.FalsingCollector
 import com.android.systemui.coroutines.collectLastValue
-import com.android.systemui.keyguard.data.repository.FakeDeviceEntryFaceAuthRepository
+import com.android.systemui.keyguard.data.repository.fakeDeviceEntryFaceAuthRepository
 import com.android.systemui.model.SysUiState
 import com.android.systemui.power.domain.interactor.PowerInteractor.Companion.setAsleepForTest
 import com.android.systemui.power.domain.interactor.PowerInteractor.Companion.setAwakeForTest
@@ -68,17 +68,11 @@
     private val utils = SceneTestUtils(this)
     private val testScope = utils.testScope
     private val sceneInteractor = utils.sceneInteractor()
-    private val sceneContainerFlags = utils.sceneContainerFlags
+    private val sceneContainerFlags = utils.fakeSceneContainerFlags
     private val authenticationInteractor = utils.authenticationInteractor()
-    private val bouncerInteractor =
-        utils.bouncerInteractor(authenticationInteractor = authenticationInteractor)
-    private val faceAuthRepository = FakeDeviceEntryFaceAuthRepository()
-    private val deviceEntryInteractor =
-        utils.deviceEntryInteractor(
-            faceAuthRepository = faceAuthRepository,
-            authenticationInteractor = authenticationInteractor,
-            sceneInteractor = sceneInteractor,
-        )
+    private val bouncerInteractor = utils.bouncerInteractor()
+    private val faceAuthRepository = utils.kosmos.fakeDeviceEntryFaceAuthRepository
+    private val deviceEntryInteractor = utils.deviceEntryInteractor()
     private val keyguardInteractor = utils.keyguardInteractor()
     private val sysUiState: SysUiState = mock()
     private val falsingCollector: FalsingCollector = mock()
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModelTest.kt
index c89cd9e..3a4ee64 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModelTest.kt
@@ -28,6 +28,7 @@
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.runTest
+import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 
@@ -37,11 +38,17 @@
 
     private val utils = SceneTestUtils(this)
     private val interactor = utils.sceneInteractor()
-    private val underTest =
-        SceneContainerViewModel(
-            sceneInteractor = interactor,
-            falsingInteractor = utils.falsingInteractor(),
-        )
+    private lateinit var underTest: SceneContainerViewModel
+
+    @Before
+    fun setUp() {
+        utils.fakeSceneContainerFlags.enabled = true
+        underTest =
+            SceneContainerViewModel(
+                sceneInteractor = interactor,
+                falsingInteractor = utils.falsingInteractor(),
+            )
+    }
 
     @Test
     fun isVisible() = runTest {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneViewModelTest.kt
index 1d2497d..a8133a3a 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneViewModelTest.kt
@@ -56,12 +56,7 @@
     private val utils = SceneTestUtils(this)
     private val testScope = utils.testScope
     private val sceneInteractor = utils.sceneInteractor()
-    private val authenticationInteractor = utils.authenticationInteractor()
-    private val deviceEntryInteractor =
-        utils.deviceEntryInteractor(
-            authenticationInteractor = authenticationInteractor,
-            sceneInteractor = sceneInteractor,
-        )
+    private val deviceEntryInteractor = utils.deviceEntryInteractor()
 
     private val mobileIconsInteractor = FakeMobileIconsInteractor(FakeMobileMappingsProxy(), mock())
     private val flags = FakeFeatureFlagsClassic().also { it.set(Flags.NEW_NETWORK_SLICE_UI, false) }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/NotificationStackAppearanceIntegrationTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/NotificationStackAppearanceIntegrationTest.kt
index 4cdb08a..607996d 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/NotificationStackAppearanceIntegrationTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/NotificationStackAppearanceIntegrationTest.kt
@@ -27,8 +27,7 @@
 import com.android.systemui.flags.fakeFeatureFlagsClassic
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.scene.domain.interactor.sceneInteractor
-import com.android.systemui.scene.shared.flag.FakeSceneContainerFlags
-import com.android.systemui.scene.shared.flag.sceneContainerFlags
+import com.android.systemui.scene.shared.flag.fakeSceneContainerFlags
 import com.android.systemui.scene.shared.model.ObservableTransitionState
 import com.android.systemui.scene.shared.model.SceneKey
 import com.android.systemui.scene.shared.model.SceneModel
@@ -50,7 +49,7 @@
 
     private val kosmos =
         testKosmos().apply {
-            sceneContainerFlags = FakeSceneContainerFlags(enabled = true)
+            fakeSceneContainerFlags.enabled = true
             fakeFeatureFlagsClassic.apply {
                 set(Flags.FULL_SCREEN_USER_SWITCHER, false)
                 set(Flags.NSSL_DEBUG_LINES, false)
diff --git a/packages/SystemUI/tests/src/com/android/keyguard/LockIconViewControllerBaseTest.java b/packages/SystemUI/tests/src/com/android/keyguard/LockIconViewControllerBaseTest.java
index 3d7d701..d8ef981 100644
--- a/packages/SystemUI/tests/src/com/android/keyguard/LockIconViewControllerBaseTest.java
+++ b/packages/SystemUI/tests/src/com/android/keyguard/LockIconViewControllerBaseTest.java
@@ -175,7 +175,7 @@
                 mPrimaryBouncerInteractor,
                 mContext,
                 () -> mDeviceEntryInteractor,
-                mSceneTestUtils.getSceneContainerFlags()
+                mSceneTestUtils.getFakeSceneContainerFlags()
         );
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/keyguard/LockIconViewControllerTest.java b/packages/SystemUI/tests/src/com/android/keyguard/LockIconViewControllerTest.java
index 93a5393..132bdb5 100644
--- a/packages/SystemUI/tests/src/com/android/keyguard/LockIconViewControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/keyguard/LockIconViewControllerTest.java
@@ -373,7 +373,7 @@
     @Test
     public void longPress_showBouncer_sceneContainerNotEnabled() {
         init(/* useMigrationFlag= */ false);
-        mSceneTestUtils.getSceneContainerFlags().setEnabled(false);
+        mSceneTestUtils.getFakeSceneContainerFlags().setEnabled(false);
         when(mFalsingManager.isFalseLongTap(anyInt())).thenReturn(false);
 
         // WHEN longPress
@@ -387,7 +387,7 @@
     @Test
     public void longPress_showBouncer() {
         init(/* useMigrationFlag= */ false);
-        mSceneTestUtils.getSceneContainerFlags().setEnabled(true);
+        mSceneTestUtils.getFakeSceneContainerFlags().setEnabled(true);
         when(mFalsingManager.isFalseLongTap(anyInt())).thenReturn(false);
 
         // WHEN longPress
@@ -401,7 +401,7 @@
     @Test
     public void longPress_falsingTriggered_doesNotShowBouncer() {
         init(/* useMigrationFlag= */ false);
-        mSceneTestUtils.getSceneContainerFlags().setEnabled(true);
+        mSceneTestUtils.getFakeSceneContainerFlags().setEnabled(true);
         when(mFalsingManager.isFalseLongTap(anyInt())).thenReturn(true);
 
         // WHEN longPress
diff --git a/packages/SystemUI/tests/src/com/android/systemui/deviceentry/data/repository/DeviceEntryHapticsInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/deviceentry/data/repository/DeviceEntryHapticsInteractorTest.kt
index d5c3641..0dfdeca 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/deviceentry/data/repository/DeviceEntryHapticsInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/deviceentry/data/repository/DeviceEntryHapticsInteractorTest.kt
@@ -36,9 +36,9 @@
 import com.android.systemui.power.shared.model.WakeSleepReason
 import com.android.systemui.power.shared.model.WakefulnessState
 import com.android.systemui.testKosmos
-import com.android.systemui.util.time.fakeSystemClock
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.advanceTimeBy
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
@@ -72,8 +72,8 @@
 
             // It's been 10 seconds since the last power button wakeup
             setAwakeFromPowerButton()
+            advanceTimeBy(10000)
             runCurrent()
-            kosmos.fakeSystemClock.setUptimeMillis(kosmos.fakeSystemClock.uptimeMillis() + 10000)
 
             enterDeviceFromBiometricUnlock()
             assertThat(playSuccessHaptic).isNotNull()
@@ -89,8 +89,8 @@
 
             // It's been 10 seconds since the last power button wakeup
             setAwakeFromPowerButton()
+            advanceTimeBy(10000)
             runCurrent()
-            kosmos.fakeSystemClock.setUptimeMillis(kosmos.fakeSystemClock.uptimeMillis() + 10000)
 
             enterDeviceFromBiometricUnlock()
             assertThat(playSuccessHaptic).isNull()
@@ -106,8 +106,8 @@
 
             // It's only been 50ms since the last power button wakeup
             setAwakeFromPowerButton()
+            advanceTimeBy(50)
             runCurrent()
-            kosmos.fakeSystemClock.setUptimeMillis(kosmos.fakeSystemClock.uptimeMillis() + 50)
 
             enterDeviceFromBiometricUnlock()
             assertThat(playSuccessHaptic).isNull()
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationShadeWindowControllerImplTest.java b/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationShadeWindowControllerImplTest.java
index 791c080..971c8a3 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationShadeWindowControllerImplTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationShadeWindowControllerImplTest.java
@@ -131,7 +131,6 @@
     @Spy private final NotificationShadeWindowView mNotificationShadeWindowView = spy(
             new NotificationShadeWindowView(mContext, null));
     @Mock private IActivityManager mActivityManager;
-    @Mock private SysuiStatusBarStateController mStatusBarStateController;
     @Mock private ConfigurationController mConfigurationController;
     @Mock private KeyguardViewMediator mKeyguardViewMediator;
     @Mock private KeyguardBypassController mKeyguardBypassController;
@@ -140,7 +139,6 @@
     @Mock private DumpManager mDumpManager;
     @Mock private KeyguardSecurityModel mKeyguardSecurityModel;
     @Mock private KeyguardStateController mKeyguardStateController;
-    @Mock private ScreenOffAnimationController mScreenOffAnimationController;
     @Mock private AuthController mAuthController;
     @Mock private ShadeWindowLogger mShadeWindowLogger;
     @Mock private SelectedUserInteractor mSelectedUserInteractor;
@@ -160,6 +158,8 @@
     private float mPreferredRefreshRate = -1;
     private FromLockscreenTransitionInteractor mFromLockscreenTransitionInteractor;
     private FromPrimaryBouncerTransitionInteractor mFromPrimaryBouncerTransitionInteractor;
+    private ScreenOffAnimationController mScreenOffAnimationController;
+    private SysuiStatusBarStateController mStatusBarStateController;
 
     @Before
     public void setUp() {
@@ -178,11 +178,9 @@
         FakeFeatureFlagsClassic featureFlags = new FakeFeatureFlagsClassic();
         FakeShadeRepository shadeRepository = new FakeShadeRepository();
 
-        PowerInteractor powerInteractor = mUtils.powerInteractor(
-                mUtils.getPowerRepository(),
-                mUtils.falsingCollector(),
-                mScreenOffAnimationController,
-                mStatusBarStateController);
+        mScreenOffAnimationController = mUtils.getScreenOffAnimationController();
+        mStatusBarStateController = spy(mUtils.getStatusBarStateController());
+        PowerInteractor powerInteractor = mUtils.powerInteractor();
 
         SceneInteractor sceneInteractor = new SceneInteractor(
                 mTestScope.getBackgroundScope(),
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerBaseTest.java b/packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerBaseTest.java
index a369f82..ca68fd8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerBaseTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerBaseTest.java
@@ -85,7 +85,6 @@
 import com.android.systemui.statusbar.NotificationShadeDepthController;
 import com.android.systemui.statusbar.PulseExpansionHandler;
 import com.android.systemui.statusbar.QsFrameTranslateController;
-import com.android.systemui.statusbar.StatusBarStateControllerImpl;
 import com.android.systemui.statusbar.SysuiStatusBarStateController;
 import com.android.systemui.statusbar.disableflags.data.repository.FakeDisableFlagsRepository;
 import com.android.systemui.statusbar.notification.data.repository.ActiveNotificationListRepository;
@@ -99,7 +98,6 @@
 import com.android.systemui.statusbar.phone.KeyguardStatusBarView;
 import com.android.systemui.statusbar.phone.LightBarController;
 import com.android.systemui.statusbar.phone.LockscreenGestureLogger;
-import com.android.systemui.statusbar.phone.ScreenOffAnimationController;
 import com.android.systemui.statusbar.phone.ScrimController;
 import com.android.systemui.statusbar.phone.StatusBarKeyguardViewManager;
 import com.android.systemui.statusbar.phone.StatusBarTouchableRegionManager;
@@ -172,7 +170,6 @@
     @Mock protected LockscreenGestureLogger mLockscreenGestureLogger;
     @Mock protected MetricsLogger mMetricsLogger;
     @Mock protected FeatureFlags mFeatureFlags;
-    @Mock protected InteractionJankMonitor mInteractionJankMonitor;
     @Mock protected ShadeLogger mShadeLogger;
     @Mock protected DumpManager mDumpManager;
     @Mock protected UiEventLogger mUiEventLogger;
@@ -186,6 +183,7 @@
     protected FakeKeyguardRepository mKeyguardRepository = new FakeKeyguardRepository();
     protected FakeShadeRepository mShadeRepository = new FakeShadeRepository();
 
+    protected InteractionJankMonitor mInteractionJankMonitor;
     protected SysuiStatusBarStateController mStatusBarStateController;
     protected ShadeInteractor mShadeInteractor;
 
@@ -205,8 +203,8 @@
     public void setup() {
         MockitoAnnotations.initMocks(this);
         when(mPanelViewControllerLazy.get()).thenReturn(mNotificationPanelViewController);
-        mStatusBarStateController = new StatusBarStateControllerImpl(mUiEventLogger,
-                mInteractionJankMonitor, mock(JavaAdapter.class), () -> mShadeInteractor);
+        mStatusBarStateController = mUtils.getStatusBarStateController();
+        mInteractionJankMonitor = mUtils.getInteractionJankMonitor();
 
         FakeDeviceProvisioningRepository deviceProvisioningRepository =
                 new FakeDeviceProvisioningRepository();
@@ -214,11 +212,7 @@
         FakeFeatureFlagsClassic featureFlags = new FakeFeatureFlagsClassic();
         FakeConfigurationRepository configurationRepository = new FakeConfigurationRepository();
 
-        PowerInteractor powerInteractor = mUtils.powerInteractor(
-                mUtils.getPowerRepository(),
-                mUtils.falsingCollector(),
-                mock(ScreenOffAnimationController.class),
-                mStatusBarStateController);
+        PowerInteractor powerInteractor = mUtils.powerInteractor();
 
         SceneInteractor sceneInteractor = new SceneInteractor(
                 mTestScope.getBackgroundScope(),
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardStatusBarViewControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardStatusBarViewControllerTest.java
index 5102b4f..5fa7f13e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardStatusBarViewControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardStatusBarViewControllerTest.java
@@ -166,7 +166,7 @@
                 mKeyguardRepository,
                 mCommandQueue,
                 PowerInteractorFactory.create().getPowerInteractor(),
-                mSceneTestUtils.getSceneContainerFlags(),
+                mSceneTestUtils.getFakeSceneContainerFlags(),
                 new FakeKeyguardBouncerRepository(),
                 new ConfigurationInteractor(new FakeConfigurationRepository()),
                 new FakeShadeRepository(),
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModelTest.kt
index 9419d63..5f9c096 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModelTest.kt
@@ -55,7 +55,7 @@
             keyguardRepository,
             mock<CommandQueue>(),
             PowerInteractorFactory.create().powerInteractor,
-            sceneTestUtils.sceneContainerFlags,
+            sceneTestUtils.fakeSceneContainerFlags,
             FakeKeyguardBouncerRepository(),
             ConfigurationInteractor(FakeConfigurationRepository()),
             FakeShadeRepository(),
diff --git a/packages/SystemUI/tests/src/com/android/systemui/user/domain/interactor/UserSwitcherInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/user/domain/interactor/UserSwitcherInteractorTest.kt
index fb5375a..e6b9d9b 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/user/domain/interactor/UserSwitcherInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/user/domain/interactor/UserSwitcherInteractorTest.kt
@@ -121,10 +121,10 @@
             SUPERVISED_USER_CREATION_APP_PACKAGE,
         )
 
-        utils.featureFlags.set(Flags.FULL_SCREEN_USER_SWITCHER, false)
+        utils.fakeFeatureFlags.set(Flags.FULL_SCREEN_USER_SWITCHER, false)
         mSetFlagsRule.enableFlags(AConfigFlags.FLAG_SWITCH_USER_ON_BG)
         spyContext = spy(context)
-        keyguardReply = KeyguardInteractorFactory.create(featureFlags = utils.featureFlags)
+        keyguardReply = KeyguardInteractorFactory.create(featureFlags = utils.fakeFeatureFlags)
         keyguardRepository = keyguardReply.repository
         userRepository = FakeUserRepository()
         refreshUsersScheduler =
@@ -363,7 +363,7 @@
     fun actions_deviceUnlocked_fullScreen() {
         createUserInteractor()
         testScope.runTest {
-            utils.featureFlags.set(Flags.FULL_SCREEN_USER_SWITCHER, true)
+            utils.fakeFeatureFlags.set(Flags.FULL_SCREEN_USER_SWITCHER, true)
             val userInfos = createUserInfos(count = 2, includeGuest = false)
 
             userRepository.setUserInfos(userInfos)
@@ -447,7 +447,7 @@
     fun actions_deviceLockedAddFromLockscreenSet_fullList_fullScreen() {
         createUserInteractor()
         testScope.runTest {
-            utils.featureFlags.set(Flags.FULL_SCREEN_USER_SWITCHER, true)
+            utils.fakeFeatureFlags.set(Flags.FULL_SCREEN_USER_SWITCHER, true)
             val userInfos = createUserInfos(count = 2, includeGuest = false)
             userRepository.setUserInfos(userInfos)
             userRepository.setSelectedUserInfo(userInfos[0])
@@ -792,7 +792,7 @@
     fun userRecordsFullScreen() {
         createUserInteractor()
         testScope.runTest {
-            utils.featureFlags.set(Flags.FULL_SCREEN_USER_SWITCHER, true)
+            utils.fakeFeatureFlags.set(Flags.FULL_SCREEN_USER_SWITCHER, true)
             val userInfos = createUserInfos(count = 3, includeGuest = false)
             userRepository.setSettings(UserSwitcherSettingsModel(isUserSwitcherEnabled = true))
             userRepository.setUserInfos(userInfos)
@@ -901,7 +901,7 @@
     fun showUserSwitcher_fullScreenEnabled_launchesFullScreenDialog() {
         createUserInteractor()
         testScope.runTest {
-            utils.featureFlags.set(Flags.FULL_SCREEN_USER_SWITCHER, true)
+            utils.fakeFeatureFlags.set(Flags.FULL_SCREEN_USER_SWITCHER, true)
 
             val expandable = mock<Expandable>()
             underTest.showUserSwitcher(expandable)
@@ -1139,7 +1139,7 @@
                         resetOrExitSessionReceiver = resetOrExitSessionReceiver,
                     ),
                 uiEventLogger = uiEventLogger,
-                featureFlags = utils.featureFlags,
+                featureFlags = utils.fakeFeatureFlags,
                 userRestrictionChecker = mock(),
             )
     }
diff --git a/packages/SystemUI/tests/utils/src/android/content/ContextKosmos.kt b/packages/SystemUI/tests/utils/src/android/content/ContextKosmos.kt
index f96c508..5e254bf 100644
--- a/packages/SystemUI/tests/utils/src/android/content/ContextKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/android/content/ContextKosmos.kt
@@ -16,9 +16,7 @@
 
 package android.content
 
-import com.android.systemui.SysuiTestableContext
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.testCase
 
-val Kosmos.testableContext: SysuiTestableContext by Kosmos.Fixture { testCase.context }
-var Kosmos.applicationContext: Context by Kosmos.Fixture { testableContext }
+var Kosmos.applicationContext: Context by Kosmos.Fixture { testCase.context }
diff --git a/packages/SystemUI/tests/utils/src/android/telephony/TelephonyManagerKosmos.kt b/packages/SystemUI/tests/utils/src/android/telephony/TelephonyManagerKosmos.kt
new file mode 100644
index 0000000..a4ee702
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/android/telephony/TelephonyManagerKosmos.kt
@@ -0,0 +1,32 @@
+/*
+ * 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 android.telephony
+
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
+import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.whenever
+import org.mockito.ArgumentMatchers.anyInt
+import org.mockito.ArgumentMatchers.anyString
+
+val Kosmos.telephonyManager by Fixture {
+    mock<TelephonyManager> {
+        whenever(createForSubscriptionId(anyInt())).thenReturn(this)
+        whenever(supplyIccLockPin(anyString()))
+            .thenReturn(PinResult(PinResult.PIN_RESULT_TYPE_SUCCESS, 3))
+    }
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/app/ActivityTaskManagerKosmos.kt
similarity index 67%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt
copy to packages/SystemUI/tests/utils/src/com/android/app/ActivityTaskManagerKosmos.kt
index f9cdc1b..fa3e8f9 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/app/ActivityTaskManagerKosmos.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023 The Android Open Source Project
+ * 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.
@@ -14,9 +14,11 @@
  * limitations under the License.
  */
 
-package com.android.systemui.scene.shared.model
+package com.android.app
 
+import android.app.ActivityTaskManager
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
+import com.android.systemui.util.mockito.mock
 
-val Kosmos.sceneContainerConfig by
-    Kosmos.Fixture { FakeSceneContainerConfigModule().sceneContainerConfig }
+val Kosmos.activityTaskManager by Fixture { mock<ActivityTaskManager>() }
diff --git a/packages/SystemUI/tests/utils/src/com/android/internal/logging/MetricsLoggerKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/internal/logging/MetricsLoggerKosmos.kt
index a1815c5..eedc0b9 100644
--- a/packages/SystemUI/tests/utils/src/com/android/internal/logging/MetricsLoggerKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/internal/logging/MetricsLoggerKosmos.kt
@@ -16,8 +16,9 @@
 
 package com.android.internal.logging
 
+import com.android.internal.logging.testing.FakeMetricsLogger
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
-import com.android.systemui.util.mockito.mock
 
-var Kosmos.metricsLogger by Fixture { mock<MetricsLogger>() }
+val Kosmos.fakeMetricsLogger by Fixture { FakeMetricsLogger() }
+val Kosmos.metricsLogger by Fixture<MetricsLogger> { fakeMetricsLogger }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/internal/util/EmergencyAffordanceManagerKosmos.kt
similarity index 68%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt
copy to packages/SystemUI/tests/utils/src/com/android/internal/util/EmergencyAffordanceManagerKosmos.kt
index f9cdc1b..3133437 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/internal/util/EmergencyAffordanceManagerKosmos.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023 The Android Open Source Project
+ * 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.
@@ -14,9 +14,10 @@
  * limitations under the License.
  */
 
-package com.android.systemui.scene.shared.model
+package com.android.internal.util
 
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
+import com.android.systemui.util.mockito.mock
 
-val Kosmos.sceneContainerConfig by
-    Kosmos.Fixture { FakeSceneContainerConfigModule().sceneContainerConfig }
+val Kosmos.emergencyAffordanceManager by Fixture { mock<EmergencyAffordanceManager>() }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/data/repository/BouncerRepositoryKosmos.kt
similarity index 65%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt
copy to packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/data/repository/BouncerRepositoryKosmos.kt
index f9cdc1b..c0f8638 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/data/repository/BouncerRepositoryKosmos.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023 The Android Open Source Project
+ * 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.
@@ -14,9 +14,14 @@
  * limitations under the License.
  */
 
-package com.android.systemui.scene.shared.model
+package com.android.systemui.bouncer.data.repository
 
+import com.android.systemui.flags.featureFlagsClassic
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
 
-val Kosmos.sceneContainerConfig by
-    Kosmos.Fixture { FakeSceneContainerConfigModule().sceneContainerConfig }
+val Kosmos.bouncerRepository by Fixture {
+    BouncerRepository(
+        flags = featureFlagsClassic,
+    )
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/data/repository/EmergencyServicesRepositoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/data/repository/EmergencyServicesRepositoryKosmos.kt
new file mode 100644
index 0000000..8851709
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/data/repository/EmergencyServicesRepositoryKosmos.kt
@@ -0,0 +1,31 @@
+/*
+ * 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.bouncer.data.repository
+
+import android.content.res.mainResources
+import com.android.systemui.common.ui.data.repository.configurationRepository
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
+import com.android.systemui.kosmos.testScope
+
+val Kosmos.emergencyServicesRepository by Fixture {
+    EmergencyServicesRepository(
+        applicationScope = testScope.backgroundScope,
+        resources = mainResources,
+        configurationRepository = configurationRepository,
+    )
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/data/repository/SimBouncerRepositoryKosmos.kt
similarity index 64%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt
copy to packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/data/repository/SimBouncerRepositoryKosmos.kt
index f9cdc1b..7af39df 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/data/repository/SimBouncerRepositoryKosmos.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023 The Android Open Source Project
+ * 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.
@@ -14,9 +14,11 @@
  * limitations under the License.
  */
 
-package com.android.systemui.scene.shared.model
+package com.android.systemui.bouncer.data.repository
 
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
 
-val Kosmos.sceneContainerConfig by
-    Kosmos.Fixture { FakeSceneContainerConfigModule().sceneContainerConfig }
+val Kosmos.fakeSimBouncerRepository by Fixture { FakeSimBouncerRepository() }
+
+val Kosmos.simBouncerRepository by Fixture<SimBouncerRepository> { fakeSimBouncerRepository }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/domain/interactor/AlternateBouncerInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/domain/interactor/AlternateBouncerInteractorKosmos.kt
index 86a4509..c4fc30d 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/domain/interactor/AlternateBouncerInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/domain/interactor/AlternateBouncerInteractorKosmos.kt
@@ -25,7 +25,7 @@
 import com.android.systemui.plugins.statusbar.statusBarStateController
 import com.android.systemui.statusbar.policy.KeyguardStateControllerImpl
 import com.android.systemui.util.mockito.mock
-import com.android.systemui.util.time.fakeSystemClock
+import com.android.systemui.util.time.systemClock
 
 var Kosmos.alternateBouncerInteractor by
     Kosmos.Fixture {
@@ -35,7 +35,7 @@
             bouncerRepository = keyguardBouncerRepository,
             fingerprintPropertyRepository = fingerprintPropertyRepository,
             biometricSettingsRepository = biometricSettingsRepository,
-            systemClock = fakeSystemClock,
+            systemClock = systemClock,
             keyguardUpdateMonitor = keyguardUpdateMonitor,
             scope = testScope.backgroundScope,
         )
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/domain/interactor/BouncerActionButtonInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/domain/interactor/BouncerActionButtonInteractorKosmos.kt
new file mode 100644
index 0000000..5ced578
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/domain/interactor/BouncerActionButtonInteractorKosmos.kt
@@ -0,0 +1,54 @@
+/*
+ * 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.bouncer.domain.interactor
+
+import android.content.Intent
+import android.content.applicationContext
+import com.android.app.activityTaskManager
+import com.android.internal.logging.metricsLogger
+import com.android.internal.util.emergencyAffordanceManager
+import com.android.systemui.authentication.domain.interactor.authenticationInteractor
+import com.android.systemui.bouncer.data.repository.emergencyServicesRepository
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.mobileConnectionsRepository
+import com.android.systemui.telephony.domain.interactor.telephonyInteractor
+import com.android.systemui.user.domain.interactor.selectedUserInteractor
+import com.android.systemui.util.mockito.mock
+import com.android.telecom.telecomManager
+
+val Kosmos.bouncerActionButtonInteractor by Fixture {
+    BouncerActionButtonInteractor(
+        applicationContext = applicationContext,
+        backgroundDispatcher = testDispatcher,
+        repository = emergencyServicesRepository,
+        mobileConnectionsRepository = mobileConnectionsRepository,
+        telephonyInteractor = telephonyInteractor,
+        authenticationInteractor = authenticationInteractor,
+        selectedUserInteractor = selectedUserInteractor,
+        activityTaskManager = activityTaskManager,
+        telecomManager = telecomManager,
+        emergencyAffordanceManager = emergencyAffordanceManager,
+        emergencyDialerIntentFactory =
+            object : EmergencyDialerIntentFactory {
+                override fun invoke(): Intent = Intent()
+            },
+        metricsLogger = metricsLogger,
+        dozeLogger = mock(),
+    )
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/domain/interactor/BouncerInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/domain/interactor/BouncerInteractorKosmos.kt
new file mode 100644
index 0000000..27803b2
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/domain/interactor/BouncerInteractorKosmos.kt
@@ -0,0 +1,40 @@
+/*
+ * 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.bouncer.domain.interactor
+
+import android.content.applicationContext
+import com.android.systemui.authentication.domain.interactor.authenticationInteractor
+import com.android.systemui.bouncer.data.repository.bouncerRepository
+import com.android.systemui.classifier.domain.interactor.falsingInteractor
+import com.android.systemui.deviceentry.domain.interactor.deviceEntryFaceAuthInteractor
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.power.domain.interactor.powerInteractor
+
+val Kosmos.bouncerInteractor by Fixture {
+    BouncerInteractor(
+        applicationScope = testScope.backgroundScope,
+        applicationContext = applicationContext,
+        repository = bouncerRepository,
+        authenticationInteractor = authenticationInteractor,
+        deviceEntryFaceAuthInteractor = deviceEntryFaceAuthInteractor,
+        falsingInteractor = falsingInteractor,
+        powerInteractor = powerInteractor,
+        simBouncerInteractor = simBouncerInteractor,
+    )
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/domain/interactor/SimBouncerInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/domain/interactor/SimBouncerInteractorKosmos.kt
new file mode 100644
index 0000000..8ed9f45
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/domain/interactor/SimBouncerInteractorKosmos.kt
@@ -0,0 +1,44 @@
+/*
+ * 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.bouncer.domain.interactor
+
+import android.content.Context
+import android.content.applicationContext
+import android.content.res.mainResources
+import android.telephony.euicc.EuiccManager
+import android.telephony.telephonyManager
+import com.android.keyguard.keyguardUpdateMonitor
+import com.android.systemui.bouncer.data.repository.simBouncerRepository
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.mobileConnectionsRepository
+
+val Kosmos.simBouncerInteractor by Fixture {
+    SimBouncerInteractor(
+        applicationContext = applicationContext,
+        backgroundDispatcher = testDispatcher,
+        applicationScope = testScope.backgroundScope,
+        repository = simBouncerRepository,
+        telephonyManager = telephonyManager,
+        resources = mainResources,
+        keyguardUpdateMonitor = keyguardUpdateMonitor,
+        euiccManager = applicationContext.getSystemService(Context.EUICC_SERVICE) as EuiccManager,
+        mobileConnectionsRepository = mobileConnectionsRepository,
+    )
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModelKosmos.kt
new file mode 100644
index 0000000..d91c597
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModelKosmos.kt
@@ -0,0 +1,49 @@
+/*
+ * 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.bouncer.ui.viewmodel
+
+import android.content.applicationContext
+import com.android.systemui.authentication.domain.interactor.authenticationInteractor
+import com.android.systemui.bouncer.domain.interactor.bouncerActionButtonInteractor
+import com.android.systemui.bouncer.domain.interactor.bouncerInteractor
+import com.android.systemui.bouncer.domain.interactor.simBouncerInteractor
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.scene.shared.flag.sceneContainerFlags
+import com.android.systemui.user.ui.viewmodel.userSwitcherViewModel
+import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.time.systemClock
+
+val Kosmos.bouncerViewModel by Fixture {
+    BouncerViewModel(
+        applicationContext = applicationContext,
+        applicationScope = testScope.backgroundScope,
+        mainDispatcher = testDispatcher,
+        bouncerInteractor = bouncerInteractor,
+        simBouncerInteractor = simBouncerInteractor,
+        authenticationInteractor = authenticationInteractor,
+        flags = sceneContainerFlags,
+        selectedUser = userSwitcherViewModel.selectedUser,
+        users = userSwitcherViewModel.users,
+        userSwitcherMenu = userSwitcherViewModel.menu,
+        actionButton = bouncerActionButtonInteractor.actionButton,
+        clock = systemClock,
+        devicePolicyManager = mock(),
+    )
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/classifier/domain/interactor/FalsingInteractorKosmos.kt
similarity index 64%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt
copy to packages/SystemUI/tests/utils/src/com/android/systemui/classifier/domain/interactor/FalsingInteractorKosmos.kt
index f9cdc1b..8fee5b2 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/classifier/domain/interactor/FalsingInteractorKosmos.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023 The Android Open Source Project
+ * 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.
@@ -14,9 +14,14 @@
  * limitations under the License.
  */
 
-package com.android.systemui.scene.shared.model
+package com.android.systemui.classifier.domain.interactor
 
+import com.android.systemui.classifier.falsingCollector
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
 
-val Kosmos.sceneContainerConfig by
-    Kosmos.Fixture { FakeSceneContainerConfigModule().sceneContainerConfig }
+val Kosmos.falsingInteractor by Fixture {
+    FalsingInteractor(
+        collector = falsingCollector,
+    )
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/CommunalMediaRepositoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/CommunalMediaRepositoryKosmos.kt
index 7946446..0768106 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/CommunalMediaRepositoryKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/CommunalMediaRepositoryKosmos.kt
@@ -18,6 +18,7 @@
 
 import com.android.systemui.kosmos.Kosmos
 
-var Kosmos.communalMediaRepository: CommunalMediaRepository by
-    Kosmos.Fixture { fakeCommunalMediaRepository }
 val Kosmos.fakeCommunalMediaRepository by Kosmos.Fixture { FakeCommunalMediaRepository() }
+
+val Kosmos.communalMediaRepository by
+    Kosmos.Fixture<CommunalMediaRepository> { fakeCommunalMediaRepository }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/CommunalRepositoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/CommunalRepositoryKosmos.kt
index be56d2b..1f5af5c 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/CommunalRepositoryKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/CommunalRepositoryKosmos.kt
@@ -17,6 +17,8 @@
 package com.android.systemui.communal.data.repository
 
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
 
-var Kosmos.communalRepository: CommunalRepository by Kosmos.Fixture { fakeCommunalRepository }
-val Kosmos.fakeCommunalRepository by Kosmos.Fixture { FakeCommunalRepository() }
+val Kosmos.fakeCommunalRepository by Fixture { FakeCommunalRepository() }
+
+val Kosmos.communalRepository by Fixture<CommunalRepository> { fakeCommunalRepository }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/CommunalWidgetRepositoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/CommunalWidgetRepositoryKosmos.kt
index 5a17f2f8..c225e3c 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/CommunalWidgetRepositoryKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/CommunalWidgetRepositoryKosmos.kt
@@ -17,9 +17,14 @@
 package com.android.systemui.communal.data.repository
 
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
 import com.android.systemui.kosmos.applicationCoroutineScope
 
-var Kosmos.communalWidgetRepository: CommunalWidgetRepository by
-    Kosmos.Fixture { fakeCommunalWidgetRepository }
-val Kosmos.fakeCommunalWidgetRepository by
-    Kosmos.Fixture { FakeCommunalWidgetRepository(applicationCoroutineScope) }
+val Kosmos.fakeCommunalWidgetRepository by Fixture {
+    FakeCommunalWidgetRepository(
+        coroutineScope = applicationCoroutineScope,
+    )
+}
+
+val Kosmos.communalWidgetRepository by
+    Fixture<CommunalWidgetRepository> { fakeCommunalWidgetRepository }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/communal/domain/interactor/CommunalInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/domain/interactor/CommunalInteractorKosmos.kt
new file mode 100644
index 0000000..649b373
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/domain/interactor/CommunalInteractorKosmos.kt
@@ -0,0 +1,38 @@
+/*
+ * 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.communal.domain.interactor
+
+import com.android.systemui.communal.data.repository.communalMediaRepository
+import com.android.systemui.communal.data.repository.communalRepository
+import com.android.systemui.communal.data.repository.communalWidgetRepository
+import com.android.systemui.keyguard.domain.interactor.keyguardInteractor
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
+import com.android.systemui.smartspace.data.repository.smartspaceRepository
+import com.android.systemui.util.mockito.mock
+
+val Kosmos.communalInteractor by Fixture {
+    CommunalInteractor(
+        communalRepository = communalRepository,
+        widgetRepository = communalWidgetRepository,
+        mediaRepository = communalMediaRepository,
+        smartspaceRepository = smartspaceRepository,
+        appWidgetHost = mock(),
+        keyguardInteractor = keyguardInteractor,
+        editWidgetsActivityStarter = mock(),
+    )
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryHapticsInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryHapticsInteractorKosmos.kt
index 6bf527d..de58ae5 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryHapticsInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryHapticsInteractorKosmos.kt
@@ -24,7 +24,7 @@
 import com.android.systemui.keyguard.data.repository.biometricSettingsRepository
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.power.domain.interactor.powerInteractor
-import com.android.systemui.util.time.fakeSystemClock
+import com.android.systemui.util.time.systemClock
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 
 val Kosmos.deviceEntryHapticsInteractor by
@@ -37,7 +37,7 @@
             biometricSettingsRepository = biometricSettingsRepository,
             keyEventInteractor = keyEventInteractor,
             powerInteractor = powerInteractor,
-            systemClock = fakeSystemClock,
+            systemClock = systemClock,
             logger = biometricUnlockLogger,
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/flags/FeatureFlagsClassicKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/flags/FeatureFlagsClassicKosmos.kt
index abadaf7..7b36a29 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/flags/FeatureFlagsClassicKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/flags/FeatureFlagsClassicKosmos.kt
@@ -30,7 +30,13 @@
  * Fixture supplying a shared [FakeFeatureFlagsClassic] instance. Can be accessed in tests in order
  * to override flag values.
  */
-val Kosmos.fakeFeatureFlagsClassic by Kosmos.Fixture { FakeFeatureFlagsClassic() }
+val Kosmos.fakeFeatureFlagsClassic by
+    Kosmos.Fixture {
+        FakeFeatureFlagsClassic().apply {
+            set(Flags.FULL_SCREEN_USER_SWITCHER, false)
+            set(Flags.NSSL_DEBUG_LINES, false)
+        }
+    }
 
 /**
  * Fixture supplying a real [FeatureFlagsClassicRelease] instance, for use by tests that want to
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/jank/InteractionJankMonitorKosmos.kt
similarity index 65%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt
copy to packages/SystemUI/tests/utils/src/com/android/systemui/jank/InteractionJankMonitorKosmos.kt
index f9cdc1b..5c5016d 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/jank/InteractionJankMonitorKosmos.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023 The Android Open Source Project
+ * 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.
@@ -14,9 +14,11 @@
  * limitations under the License.
  */
 
-package com.android.systemui.scene.shared.model
+package com.android.systemui.jank
 
+import com.android.internal.jank.InteractionJankMonitor
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
+import com.android.systemui.util.mockito.mock
 
-val Kosmos.sceneContainerConfig by
-    Kosmos.Fixture { FakeSceneContainerConfigModule().sceneContainerConfig }
+val Kosmos.interactionJankMonitor by Fixture<InteractionJankMonitor> { mock() }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/plugins/statusbar/StatusBarStateControllerKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/plugins/statusbar/StatusBarStateControllerKosmos.kt
index cac2646..73b7c50 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/plugins/statusbar/StatusBarStateControllerKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/plugins/statusbar/StatusBarStateControllerKosmos.kt
@@ -16,7 +16,20 @@
 
 package com.android.systemui.plugins.statusbar
 
+import com.android.systemui.jank.interactionJankMonitor
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.qs.uiEventLogger
+import com.android.systemui.shade.domain.interactor.shadeInteractor
+import com.android.systemui.statusbar.StatusBarStateControllerImpl
 import com.android.systemui.util.mockito.mock
 
-var Kosmos.statusBarStateController by Kosmos.Fixture { mock<StatusBarStateController>() }
+var Kosmos.statusBarStateController by
+    Kosmos.Fixture {
+        StatusBarStateControllerImpl(
+            uiEventLogger,
+            interactionJankMonitor,
+            mock(),
+        ) {
+            shadeInteractor
+        }
+    }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneTestUtils.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneTestUtils.kt
index 09ab655..d314a25 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneTestUtils.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneTestUtils.kt
@@ -16,182 +16,108 @@
 
 package com.android.systemui.scene
 
-import android.app.ActivityTaskManager
-import android.app.admin.DevicePolicyManager
 import android.content.Context
-import android.content.Intent
-import android.content.pm.UserInfo
-import android.graphics.Bitmap
-import android.graphics.drawable.BitmapDrawable
-import android.telecom.TelecomManager
-import android.telephony.PinResult
-import android.telephony.PinResult.PIN_RESULT_TYPE_SUCCESS
-import android.telephony.TelephonyManager
-import android.telephony.euicc.EuiccManager
-import com.android.internal.logging.MetricsLogger
-import com.android.internal.util.EmergencyAffordanceManager
+import android.content.applicationContext
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.authentication.data.repository.AuthenticationRepository
-import com.android.systemui.authentication.data.repository.FakeAuthenticationRepository
+import com.android.systemui.authentication.data.repository.fakeAuthenticationRepository
 import com.android.systemui.authentication.domain.interactor.AuthenticationInteractor
-import com.android.systemui.bouncer.data.repository.BouncerRepository
-import com.android.systemui.bouncer.data.repository.EmergencyServicesRepository
-import com.android.systemui.bouncer.data.repository.FakeKeyguardBouncerRepository
-import com.android.systemui.bouncer.data.repository.FakeSimBouncerRepository
+import com.android.systemui.authentication.domain.interactor.authenticationInteractor
+import com.android.systemui.bouncer.data.repository.bouncerRepository
+import com.android.systemui.bouncer.data.repository.fakeSimBouncerRepository
 import com.android.systemui.bouncer.domain.interactor.BouncerActionButtonInteractor
 import com.android.systemui.bouncer.domain.interactor.BouncerInteractor
-import com.android.systemui.bouncer.domain.interactor.EmergencyDialerIntentFactory
-import com.android.systemui.bouncer.domain.interactor.SimBouncerInteractor
+import com.android.systemui.bouncer.domain.interactor.bouncerActionButtonInteractor
+import com.android.systemui.bouncer.domain.interactor.bouncerInteractor
+import com.android.systemui.bouncer.domain.interactor.simBouncerInteractor
 import com.android.systemui.bouncer.ui.viewmodel.BouncerViewModel
+import com.android.systemui.bouncer.ui.viewmodel.bouncerViewModel
 import com.android.systemui.classifier.FalsingCollector
-import com.android.systemui.classifier.FalsingCollectorFake
 import com.android.systemui.classifier.domain.interactor.FalsingInteractor
-import com.android.systemui.common.shared.model.Text
-import com.android.systemui.common.ui.data.repository.FakeConfigurationRepository
-import com.android.systemui.common.ui.domain.interactor.ConfigurationInteractor
-import com.android.systemui.communal.data.repository.FakeCommunalRepository
+import com.android.systemui.classifier.domain.interactor.falsingInteractor
+import com.android.systemui.classifier.falsingCollector
+import com.android.systemui.common.ui.data.repository.fakeConfigurationRepository
+import com.android.systemui.common.ui.domain.interactor.configurationInteractor
+import com.android.systemui.communal.data.repository.fakeCommunalRepository
 import com.android.systemui.communal.domain.interactor.CommunalInteractor
-import com.android.systemui.communal.domain.interactor.CommunalInteractorFactory
-import com.android.systemui.deviceentry.data.repository.DeviceEntryFaceAuthRepository
-import com.android.systemui.deviceentry.data.repository.DeviceEntryRepository
-import com.android.systemui.deviceentry.data.repository.FakeDeviceEntryRepository
-import com.android.systemui.deviceentry.domain.interactor.DeviceEntryFaceAuthInteractor
+import com.android.systemui.communal.domain.interactor.communalInteractor
+import com.android.systemui.deviceentry.data.repository.fakeDeviceEntryRepository
 import com.android.systemui.deviceentry.domain.interactor.DeviceEntryInteractor
-import com.android.systemui.doze.DozeLogger
-import com.android.systemui.flags.FakeFeatureFlagsClassic
-import com.android.systemui.flags.Flags
-import com.android.systemui.keyguard.data.repository.FakeCommandQueue
-import com.android.systemui.keyguard.data.repository.FakeDeviceEntryFaceAuthRepository
-import com.android.systemui.keyguard.data.repository.FakeKeyguardRepository
-import com.android.systemui.keyguard.data.repository.FakeTrustRepository
-import com.android.systemui.keyguard.data.repository.KeyguardRepository
-import com.android.systemui.keyguard.data.repository.TrustRepository
+import com.android.systemui.deviceentry.domain.interactor.deviceEntryInteractor
+import com.android.systemui.flags.fakeFeatureFlagsClassic
+import com.android.systemui.jank.interactionJankMonitor
+import com.android.systemui.keyguard.data.repository.fakeKeyguardRepository
 import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor
-import com.android.systemui.keyguard.ui.viewmodel.KeyguardRootViewModel
+import com.android.systemui.keyguard.domain.interactor.keyguardInteractor
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.testCase
 import com.android.systemui.kosmos.testDispatcher
 import com.android.systemui.kosmos.testScope
-import com.android.systemui.plugins.statusbar.StatusBarStateController
-import com.android.systemui.power.data.repository.FakePowerRepository
-import com.android.systemui.power.data.repository.PowerRepository
+import com.android.systemui.plugins.statusbar.statusBarStateController
+import com.android.systemui.power.data.repository.fakePowerRepository
 import com.android.systemui.power.domain.interactor.PowerInteractor
-import com.android.systemui.power.domain.interactor.PowerInteractorFactory
+import com.android.systemui.power.domain.interactor.powerInteractor
 import com.android.systemui.scene.data.repository.SceneContainerRepository
+import com.android.systemui.scene.data.repository.sceneContainerRepository
 import com.android.systemui.scene.domain.interactor.SceneInteractor
-import com.android.systemui.scene.shared.flag.FakeSceneContainerFlags
+import com.android.systemui.scene.domain.interactor.sceneInteractor
+import com.android.systemui.scene.shared.flag.fakeSceneContainerFlags
 import com.android.systemui.scene.shared.model.SceneContainerConfig
 import com.android.systemui.scene.shared.model.SceneKey
-import com.android.systemui.shade.data.repository.FakeShadeRepository
-import com.android.systemui.statusbar.notification.stack.data.repository.NotificationStackAppearanceRepository
-import com.android.systemui.statusbar.notification.stack.domain.interactor.NotificationStackAppearanceInteractor
 import com.android.systemui.statusbar.notification.stack.ui.viewmodel.NotificationsPlaceholderViewModel
-import com.android.systemui.statusbar.phone.ScreenOffAnimationController
-import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionsRepository
-import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileConnectionsRepository
-import com.android.systemui.telephony.data.repository.FakeTelephonyRepository
-import com.android.systemui.telephony.data.repository.TelephonyRepository
+import com.android.systemui.statusbar.notification.stack.ui.viewmodel.notificationsPlaceholderViewModel
+import com.android.systemui.statusbar.phone.screenOffAnimationController
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.fakeMobileConnectionsRepository
+import com.android.systemui.telephony.data.repository.fakeTelephonyRepository
 import com.android.systemui.telephony.domain.interactor.TelephonyInteractor
-import com.android.systemui.user.data.repository.FakeUserRepository
+import com.android.systemui.telephony.domain.interactor.telephonyInteractor
 import com.android.systemui.user.domain.interactor.SelectedUserInteractor
-import com.android.systemui.user.ui.viewmodel.UserActionViewModel
-import com.android.systemui.user.ui.viewmodel.UserViewModel
-import com.android.systemui.util.mockito.mock
-import com.android.systemui.util.mockito.whenever
-import com.android.systemui.util.time.SystemClock
-import kotlinx.coroutines.CoroutineScope
+import com.android.systemui.user.domain.interactor.selectedUserInteractor
+import com.android.systemui.util.time.systemClock
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.flow.flowOf
-import kotlinx.coroutines.runBlocking
-import kotlinx.coroutines.test.currentTime
-import org.mockito.ArgumentMatchers.anyInt
-import org.mockito.ArgumentMatchers.anyString
 
 /**
  * Utilities for creating scene container framework related repositories, interactors, and
  * view-models for tests.
  */
 @OptIn(ExperimentalCoroutinesApi::class)
-class SceneTestUtils(
-    private val context: Context,
-) {
-    constructor(test: SysuiTestCase) : this(context = test.context)
+@Deprecated("Please use Kosmos instead.")
+class SceneTestUtils {
 
     val kosmos = Kosmos()
-    val testDispatcher = kosmos.testDispatcher
-    val testScope = kosmos.testScope
-    val featureFlags =
-        FakeFeatureFlagsClassic().apply {
-            set(Flags.FULL_SCREEN_USER_SWITCHER, false)
-            set(Flags.NSSL_DEBUG_LINES, false)
-        }
-    val sceneContainerFlags = FakeSceneContainerFlags().apply { enabled = true }
-    val deviceEntryRepository: FakeDeviceEntryRepository by lazy { FakeDeviceEntryRepository() }
-    val authenticationRepository: FakeAuthenticationRepository by lazy {
-        FakeAuthenticationRepository(
-            currentTime = { testScope.currentTime },
-        )
-    }
-    val configurationRepository: FakeConfigurationRepository by lazy {
-        FakeConfigurationRepository()
-    }
-    val configurationInteractor: ConfigurationInteractor by lazy {
-        ConfigurationInteractor(configurationRepository)
-    }
-    private val emergencyServicesRepository: EmergencyServicesRepository by lazy {
-        EmergencyServicesRepository(
-            applicationScope = applicationScope(),
-            resources = context.resources,
-            configurationRepository = configurationRepository,
-        )
-    }
-    val telephonyRepository: FakeTelephonyRepository by lazy { FakeTelephonyRepository() }
-    val bouncerRepository = BouncerRepository(featureFlags)
-    val communalRepository: FakeCommunalRepository by lazy { FakeCommunalRepository() }
-    val keyguardRepository: FakeKeyguardRepository by lazy { FakeKeyguardRepository() }
-    val powerRepository: FakePowerRepository by lazy { FakePowerRepository() }
-    val simBouncerRepository: FakeSimBouncerRepository by lazy { FakeSimBouncerRepository() }
 
-    val clock: SystemClock = mock {
-        whenever(elapsedRealtime()).thenAnswer { testScope.currentTime }
-    }
-    val telephonyManager: TelephonyManager = mock {
-        whenever(createForSubscriptionId(anyInt())).thenReturn(this)
-        whenever(supplyIccLockPin(anyString())).thenReturn(PinResult(PIN_RESULT_TYPE_SUCCESS, 3))
-    }
-    val devicePolicyManager: DevicePolicyManager = mock {}
-    val mobileConnectionsRepository: FakeMobileConnectionsRepository by lazy {
-        FakeMobileConnectionsRepository(mock(), mock())
+    constructor(
+        context: Context,
+    ) {
+        kosmos.applicationContext = context
     }
 
-    val simBouncerInteractor =
-        SimBouncerInteractor(
-            applicationContext = context,
-            backgroundDispatcher = testDispatcher,
-            applicationScope = applicationScope(),
-            repository = simBouncerRepository,
-            telephonyManager = telephonyManager,
-            resources = context.resources,
-            keyguardUpdateMonitor = mock(),
-            euiccManager = context.getSystemService(Context.EUICC_SERVICE) as EuiccManager,
-            mobileConnectionsRepository = mobileConnectionsRepository,
-        )
-
-    val userRepository: FakeUserRepository by lazy {
-        FakeUserRepository().apply {
-            val users = listOf(UserInfo(/* id=  */ 0, "name", /* flags= */ 0))
-            setUserInfos(users)
-            runBlocking { setSelectedUserInfo(users.first()) }
-        }
+    constructor(testCase: SysuiTestCase) : this(context = testCase.context) {
+        kosmos.testCase = testCase
     }
 
-    private val falsingCollectorFake: FalsingCollector by lazy { FalsingCollectorFake() }
-    private var falsingInteractor: FalsingInteractor? = null
-    private var powerInteractor: PowerInteractor? = null
+    val testDispatcher by lazy { kosmos.testDispatcher }
+    val testScope by lazy { kosmos.testScope }
+    val fakeFeatureFlags by lazy { kosmos.fakeFeatureFlagsClassic }
+    val fakeSceneContainerFlags by lazy { kosmos.fakeSceneContainerFlags }
+    val deviceEntryRepository by lazy { kosmos.fakeDeviceEntryRepository }
+    val authenticationRepository by lazy { kosmos.fakeAuthenticationRepository }
+    val configurationRepository by lazy { kosmos.fakeConfigurationRepository }
+    val configurationInteractor by lazy { kosmos.configurationInteractor }
+    val telephonyRepository by lazy { kosmos.fakeTelephonyRepository }
+    val bouncerRepository by lazy { kosmos.bouncerRepository }
+    val communalRepository by lazy { kosmos.fakeCommunalRepository }
+    val keyguardRepository by lazy { kosmos.fakeKeyguardRepository }
+    val powerRepository by lazy { kosmos.fakePowerRepository }
+    val simBouncerRepository by lazy { kosmos.fakeSimBouncerRepository }
+    val clock by lazy { kosmos.systemClock }
+    val mobileConnectionsRepository by lazy { kosmos.fakeMobileConnectionsRepository }
+    val simBouncerInteractor by lazy { kosmos.simBouncerInteractor }
+    val statusBarStateController by lazy { kosmos.statusBarStateController }
+    val interactionJankMonitor by lazy { kosmos.interactionJankMonitor }
+    val screenOffAnimationController by lazy { kosmos.screenOffAnimationController }
 
-    fun fakeSceneContainerRepository(
-        containerConfig: SceneContainerConfig = fakeSceneContainerConfig(),
-    ): SceneContainerRepository {
-        return SceneContainerRepository(applicationScope(), containerConfig)
+    fun fakeSceneContainerRepository(): SceneContainerRepository {
+        return kosmos.sceneContainerRepository
     }
 
     fun fakeSceneKeys(): List<SceneKey> {
@@ -202,222 +128,59 @@
         return kosmos.sceneContainerConfig
     }
 
-    @JvmOverloads
-    fun sceneInteractor(
-        repository: SceneContainerRepository = fakeSceneContainerRepository()
-    ): SceneInteractor {
-        return SceneInteractor(
-            applicationScope = applicationScope(),
-            repository = repository,
-            powerInteractor = powerInteractor(),
-            logger = mock(),
-        )
+    fun sceneInteractor(): SceneInteractor {
+        return kosmos.sceneInteractor
     }
 
-    fun deviceEntryInteractor(
-        repository: DeviceEntryRepository = deviceEntryRepository,
-        authenticationInteractor: AuthenticationInteractor,
-        sceneInteractor: SceneInteractor,
-        faceAuthRepository: DeviceEntryFaceAuthRepository = FakeDeviceEntryFaceAuthRepository(),
-        trustRepository: TrustRepository = FakeTrustRepository(),
-    ): DeviceEntryInteractor {
-        return DeviceEntryInteractor(
-            applicationScope = applicationScope(),
-            repository = repository,
-            authenticationInteractor = authenticationInteractor,
-            sceneInteractor = sceneInteractor,
-            deviceEntryFaceAuthRepository = faceAuthRepository,
-            trustRepository = trustRepository,
-            flags = FakeSceneContainerFlags(enabled = true)
-        )
+    fun deviceEntryInteractor(): DeviceEntryInteractor {
+        return kosmos.deviceEntryInteractor
     }
 
-    fun authenticationInteractor(
-        repository: AuthenticationRepository = authenticationRepository,
-    ): AuthenticationInteractor {
-        return AuthenticationInteractor(
-            applicationScope = applicationScope(),
-            repository = repository,
-            selectedUserInteractor = selectedUserInteractor(),
-        )
+    fun authenticationInteractor(): AuthenticationInteractor {
+        return kosmos.authenticationInteractor
     }
 
-    fun keyguardInteractor(
-        repository: KeyguardRepository = keyguardRepository
-    ): KeyguardInteractor {
-        return KeyguardInteractor(
-            repository = repository,
-            commandQueue = FakeCommandQueue(),
-            sceneContainerFlags = sceneContainerFlags,
-            bouncerRepository = FakeKeyguardBouncerRepository(),
-            configurationInteractor = configurationInteractor,
-            shadeRepository = FakeShadeRepository(),
-            sceneInteractorProvider = { sceneInteractor() },
-            powerInteractor = PowerInteractorFactory.create().powerInteractor,
-        )
+    fun keyguardInteractor(): KeyguardInteractor {
+        return kosmos.keyguardInteractor
     }
 
     fun communalInteractor(): CommunalInteractor {
-        return CommunalInteractorFactory.create(
-                communalRepository = communalRepository,
-            )
-            .communalInteractor
+        return kosmos.communalInteractor
     }
 
-    fun bouncerInteractor(
-        authenticationInteractor: AuthenticationInteractor,
-        deviceEntryFaceAuthInteractor: DeviceEntryFaceAuthInteractor = mock(),
-    ): BouncerInteractor {
-        return BouncerInteractor(
-            applicationScope = applicationScope(),
-            applicationContext = context,
-            repository = bouncerRepository,
-            authenticationInteractor = authenticationInteractor,
-            deviceEntryFaceAuthInteractor = deviceEntryFaceAuthInteractor,
-            falsingInteractor = falsingInteractor(),
-            powerInteractor = powerInteractor(),
-            simBouncerInteractor = simBouncerInteractor,
-        )
+    fun bouncerInteractor(): BouncerInteractor {
+        return kosmos.bouncerInteractor
     }
 
-    fun keyguardRootViewModel(): KeyguardRootViewModel = mock()
-
     fun notificationsPlaceholderViewModel(): NotificationsPlaceholderViewModel {
-        return NotificationsPlaceholderViewModel(
-            interactor =
-                NotificationStackAppearanceInteractor(
-                    repository = NotificationStackAppearanceRepository(),
-                ),
-            flags = sceneContainerFlags,
-            featureFlags = featureFlags,
-        )
+        return kosmos.notificationsPlaceholderViewModel
     }
 
-    fun bouncerViewModel(
-        bouncerInteractor: BouncerInteractor,
-        authenticationInteractor: AuthenticationInteractor,
-        actionButtonInteractor: BouncerActionButtonInteractor = bouncerActionButtonInteractor(),
-        users: List<UserViewModel> = createUsers(),
-    ): BouncerViewModel {
-        return BouncerViewModel(
-            applicationContext = context,
-            applicationScope = applicationScope(),
-            mainDispatcher = testDispatcher,
-            bouncerInteractor = bouncerInteractor,
-            simBouncerInteractor = simBouncerInteractor,
-            authenticationInteractor = authenticationInteractor,
-            flags = sceneContainerFlags,
-            selectedUser = flowOf(users.first { it.isSelectionMarkerVisible }),
-            users = flowOf(users),
-            userSwitcherMenu = flowOf(createMenuActions()),
-            actionButton = actionButtonInteractor.actionButton,
-            clock = clock,
-            devicePolicyManager = devicePolicyManager,
-        )
+    fun bouncerViewModel(): BouncerViewModel {
+        return kosmos.bouncerViewModel
     }
 
-    fun telephonyInteractor(
-        repository: TelephonyRepository = telephonyRepository,
-    ): TelephonyInteractor {
-        return TelephonyInteractor(repository = repository)
+    fun telephonyInteractor(): TelephonyInteractor {
+        return kosmos.telephonyInteractor
     }
 
-    fun falsingInteractor(collector: FalsingCollector = falsingCollector()): FalsingInteractor {
-        return falsingInteractor ?: FalsingInteractor(collector).also { falsingInteractor = it }
+    fun falsingInteractor(): FalsingInteractor {
+        return kosmos.falsingInteractor
     }
 
     fun falsingCollector(): FalsingCollector {
-        return falsingCollectorFake
+        return kosmos.falsingCollector
     }
 
-    fun powerInteractor(
-        repository: PowerRepository = powerRepository,
-        falsingCollector: FalsingCollector = falsingCollector(),
-        screenOffAnimationController: ScreenOffAnimationController = mock(),
-        statusBarStateController: StatusBarStateController = mock(),
-    ): PowerInteractor {
-        return powerInteractor
-            ?: PowerInteractor(
-                    repository = repository,
-                    falsingCollector = falsingCollector,
-                    screenOffAnimationController = screenOffAnimationController,
-                    statusBarStateController = statusBarStateController,
-                )
-                .also { powerInteractor = it }
-    }
-
-    private fun applicationScope(): CoroutineScope {
-        return testScope.backgroundScope
-    }
-
-    private fun createUsers(
-        count: Int = 3,
-        selectedIndex: Int = 0,
-    ): List<UserViewModel> {
-        check(selectedIndex in 0 until count)
-
-        return buildList {
-            repeat(count) { index ->
-                add(
-                    UserViewModel(
-                        viewKey = index,
-                        name = Text.Loaded("name_$index"),
-                        image = BitmapDrawable(Bitmap.createBitmap(1, 1, Bitmap.Config.ARGB_8888)),
-                        isSelectionMarkerVisible = index == selectedIndex,
-                        alpha = 1f,
-                        onClicked = {},
-                    )
-                )
-            }
-        }
-    }
-
-    private fun createMenuActions(): List<UserActionViewModel> {
-        return buildList {
-            repeat(3) { index ->
-                add(
-                    UserActionViewModel(
-                        viewKey = index.toLong(),
-                        iconResourceId = 0,
-                        textResourceId = 0,
-                        onClicked = {},
-                    )
-                )
-            }
-        }
+    fun powerInteractor(): PowerInteractor {
+        return kosmos.powerInteractor
     }
 
     fun selectedUserInteractor(): SelectedUserInteractor {
-        return SelectedUserInteractor(userRepository)
+        return kosmos.selectedUserInteractor
     }
 
-    fun bouncerActionButtonInteractor(
-        mobileConnectionsRepository: MobileConnectionsRepository = mock(),
-        activityTaskManager: ActivityTaskManager = mock(),
-        telecomManager: TelecomManager? = null,
-        emergencyAffordanceManager: EmergencyAffordanceManager =
-            EmergencyAffordanceManager(context),
-        emergencyDialerIntentFactory: EmergencyDialerIntentFactory =
-            object : EmergencyDialerIntentFactory {
-                override fun invoke(): Intent = Intent()
-            },
-        metricsLogger: MetricsLogger = mock(),
-        dozeLogger: DozeLogger = mock(),
-    ): BouncerActionButtonInteractor {
-        return BouncerActionButtonInteractor(
-            applicationContext = context,
-            backgroundDispatcher = testDispatcher,
-            repository = emergencyServicesRepository,
-            mobileConnectionsRepository = mobileConnectionsRepository,
-            telephonyInteractor = telephonyInteractor(),
-            authenticationInteractor = authenticationInteractor(),
-            selectedUserInteractor = selectedUserInteractor(),
-            activityTaskManager = activityTaskManager,
-            telecomManager = telecomManager,
-            emergencyAffordanceManager = emergencyAffordanceManager,
-            emergencyDialerIntentFactory = emergencyDialerIntentFactory,
-            metricsLogger = metricsLogger,
-            dozeLogger = dozeLogger,
-        )
+    fun bouncerActionButtonInteractor(): BouncerActionButtonInteractor {
+        return kosmos.bouncerActionButtonInteractor
     }
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryKosmos.kt
index 7c4e160..e19941c 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryKosmos.kt
@@ -18,7 +18,7 @@
 
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.applicationCoroutineScope
-import com.android.systemui.scene.shared.model.sceneContainerConfig
+import com.android.systemui.scene.sceneContainerConfig
 
 val Kosmos.sceneContainerRepository by
     Kosmos.Fixture { SceneContainerRepository(applicationCoroutineScope, sceneContainerConfig) }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/flag/SceneContainerFlagsKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/flag/SceneContainerFlagsKosmos.kt
index c2cdbed..979d8e7 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/flag/SceneContainerFlagsKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/flag/SceneContainerFlagsKosmos.kt
@@ -18,4 +18,5 @@
 
 import com.android.systemui.kosmos.Kosmos
 
-var Kosmos.sceneContainerFlags by Kosmos.Fixture { FakeSceneContainerFlags() }
+var Kosmos.fakeSceneContainerFlags by Kosmos.Fixture { FakeSceneContainerFlags() }
+val Kosmos.sceneContainerFlags by Kosmos.Fixture<SceneContainerFlags> { fakeSceneContainerFlags }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/FakeSceneContainerConfigModule.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/FakeSceneContainerConfigModule.kt
index b4fc948..8811b8d 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/FakeSceneContainerConfigModule.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/FakeSceneContainerConfigModule.kt
@@ -30,6 +30,7 @@
                     SceneKey.Lockscreen,
                     SceneKey.Bouncer,
                     SceneKey.Gone,
+                    SceneKey.Communal,
                 ),
             initialSceneKey = SceneKey.Lockscreen,
         ),
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/smartspace/data/repository/SmartspaceRepositoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/smartspace/data/repository/SmartspaceRepositoryKosmos.kt
index e671d45..0e4c923 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/smartspace/data/repository/SmartspaceRepositoryKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/smartspace/data/repository/SmartspaceRepositoryKosmos.kt
@@ -17,6 +17,8 @@
 package com.android.systemui.smartspace.data.repository
 
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
 
-var Kosmos.smartspaceRepository: SmartspaceRepository by Kosmos.Fixture { fakeSmartspaceRepository }
-val Kosmos.fakeSmartspaceRepository by Kosmos.Fixture { FakeSmartspaceRepository() }
+val Kosmos.fakeSmartspaceRepository by Fixture { FakeSmartspaceRepository() }
+
+val Kosmos.smartspaceRepository by Fixture<SmartspaceRepository> { fakeSmartspaceRepository }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/LockscreenShadeScrimTransitionControllerKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/LockscreenShadeScrimTransitionControllerKosmos.kt
index 93a7adf..8385403 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/LockscreenShadeScrimTransitionControllerKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/LockscreenShadeScrimTransitionControllerKosmos.kt
@@ -16,7 +16,7 @@
 
 package com.android.systemui.statusbar
 
-import android.content.testableContext
+import android.content.applicationContext
 import com.android.systemui.dump.dumpManager
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
@@ -27,7 +27,7 @@
 val Kosmos.lockscreenShadeScrimTransitionController by Fixture {
     LockscreenShadeScrimTransitionController(
         scrimController = scrimController,
-        context = testableContext,
+        context = applicationContext,
         configurationController = configurationController,
         dumpManager = dumpManager,
         splitShadeStateController = splitShadeStateController,
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/LockscreenShadeTransitionControllerKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/LockscreenShadeTransitionControllerKosmos.kt
index 2752cc2..1c6ce79 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/LockscreenShadeTransitionControllerKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/LockscreenShadeTransitionControllerKosmos.kt
@@ -16,7 +16,7 @@
 
 package com.android.systemui.statusbar
 
-import android.content.testableContext
+import android.content.applicationContext
 import com.android.systemui.classifier.falsingCollector
 import com.android.systemui.classifier.falsingManager
 import com.android.systemui.dump.dumpManager
@@ -47,7 +47,7 @@
         scrimTransitionController = lockscreenShadeScrimTransitionController,
         keyguardTransitionControllerFactory = lockscreenShadeKeyguardTransitionControllerFactory,
         depthController = notificationShadeDepthController,
-        context = testableContext,
+        context = applicationContext,
         splitShadeOverScrollerFactory = splitShadeLockScreenOverScrollerFactory,
         singleShadeOverScrollerFactory = singleShadeLockScreenOverScrollerFactory,
         activityStarter = activityStarter,
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/stack/AmbientStateKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/stack/AmbientStateKosmos.kt
index 83ac330..7f6f698 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/stack/AmbientStateKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/stack/AmbientStateKosmos.kt
@@ -16,7 +16,7 @@
 
 package com.android.systemui.statusbar.notification.stack
 
-import android.content.testableContext
+import android.content.applicationContext
 import com.android.systemui.dump.dumpManager
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
@@ -27,7 +27,7 @@
 @OptIn(ExperimentalCoroutinesApi::class)
 val Kosmos.ambientState by Fixture {
     AmbientState(
-        /*context=*/ testableContext,
+        /*context=*/ applicationContext,
         /*dumpManager=*/ dumpManager,
         /*sectionProvider=*/ stackScrollAlgorithmSectionProvider,
         /*bypassController=*/ stackScrollAlgorithmBypassController,
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepositoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepositoryKosmos.kt
new file mode 100644
index 0000000..9d62ea5
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepositoryKosmos.kt
@@ -0,0 +1,28 @@
+/*
+ * 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.statusbar.pipeline.mobile.data.repository
+
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
+import com.android.systemui.util.mockito.mock
+
+val Kosmos.fakeMobileConnectionsRepository by Fixture {
+    FakeMobileConnectionsRepository(tableLogBuffer = mock())
+}
+
+val Kosmos.mobileConnectionsRepository by
+    Fixture<MobileConnectionsRepository> { fakeMobileConnectionsRepository }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/user/ui/viewmodel/UserSwitcherViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/user/ui/viewmodel/UserSwitcherViewModelKosmos.kt
new file mode 100644
index 0000000..0b9f897
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/user/ui/viewmodel/UserSwitcherViewModelKosmos.kt
@@ -0,0 +1,29 @@
+/*
+ * 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.user.ui.viewmodel
+
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
+import com.android.systemui.user.domain.interactor.guestUserInteractor
+import com.android.systemui.user.domain.interactor.userSwitcherInteractor
+
+val Kosmos.userSwitcherViewModel by Fixture {
+    UserSwitcherViewModel(
+        userSwitcherInteractor = userSwitcherInteractor,
+        guestUserInteractor = guestUserInteractor,
+    )
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/util/time/FakeSystemClockKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/util/time/FakeSystemClockKosmos.kt
index 914e654..f3a8b14 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/util/time/FakeSystemClockKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/util/time/FakeSystemClockKosmos.kt
@@ -17,5 +17,19 @@
 package com.android.systemui.util.time
 
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.whenever
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.currentTime
 
-var Kosmos.fakeSystemClock by Kosmos.Fixture { FakeSystemClock() }
+@OptIn(ExperimentalCoroutinesApi::class)
+val Kosmos.systemClock by
+    Kosmos.Fixture<SystemClock> {
+        mock {
+            whenever(elapsedRealtime()).thenAnswer { testScope.currentTime }
+            whenever(uptimeMillis()).thenAnswer { testScope.currentTime }
+        }
+    }
+
+val Kosmos.fakeSystemClock by Kosmos.Fixture { FakeSystemClock() }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/telecom/TelecomManagerKosmos.kt
similarity index 68%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt
copy to packages/SystemUI/tests/utils/src/com/android/telecom/TelecomManagerKosmos.kt
index f9cdc1b..4e0c0883 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/SceneContainerConfigKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/telecom/TelecomManagerKosmos.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023 The Android Open Source Project
+ * 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.
@@ -14,9 +14,11 @@
  * limitations under the License.
  */
 
-package com.android.systemui.scene.shared.model
+package com.android.telecom
 
+import android.telecom.TelecomManager
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
+import com.android.systemui.util.mockito.mock
 
-val Kosmos.sceneContainerConfig by
-    Kosmos.Fixture { FakeSceneContainerConfigModule().sceneContainerConfig }
+var Kosmos.telecomManager by Fixture<TelecomManager?> { mock() }
diff --git a/services/companion/java/com/android/server/companion/virtual/InputController.java b/services/companion/java/com/android/server/companion/virtual/InputController.java
index d1274d4..980ec02 100644
--- a/services/companion/java/com/android/server/companion/virtual/InputController.java
+++ b/services/companion/java/com/android/server/companion/virtual/InputController.java
@@ -437,7 +437,7 @@
         }
     }
 
-    @VisibleForTesting
+    @VisibleForTesting(visibility = VisibleForTesting.Visibility.PACKAGE)
     Map<IBinder, InputDeviceDescriptor> getInputDeviceDescriptors() {
         final Map<IBinder, InputDeviceDescriptor> inputDeviceDescriptors = new ArrayMap<>();
         synchronized (mLock) {
@@ -529,7 +529,8 @@
         }
     }
 
-    @VisibleForTesting static final class InputDeviceDescriptor {
+    @VisibleForTesting(visibility = VisibleForTesting.Visibility.PACKAGE)
+    static final class InputDeviceDescriptor {
 
         static final int TYPE_KEYBOARD = 1;
         static final int TYPE_MOUSE = 2;
diff --git a/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java b/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java
index 44c3a8d..50eab25b 100644
--- a/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java
+++ b/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java
@@ -1337,6 +1337,11 @@
         }
     }
 
+    boolean isInputDeviceOwnedByVirtualDevice(int inputDeviceId) {
+        return mInputController.getInputDeviceDescriptors().values().stream().anyMatch(
+                inputDeviceDescriptor -> inputDeviceDescriptor.getInputDeviceId() == inputDeviceId);
+    }
+
     void onEnteringPipBlocked(int uid) {
         // Do nothing. ActivityRecord#checkEnterPictureInPictureState logs that the display does not
         // support PiP.
diff --git a/services/companion/java/com/android/server/companion/virtual/VirtualDeviceManagerService.java b/services/companion/java/com/android/server/companion/virtual/VirtualDeviceManagerService.java
index 0d5cdcb..ef61498 100644
--- a/services/companion/java/com/android/server/companion/virtual/VirtualDeviceManagerService.java
+++ b/services/companion/java/com/android/server/companion/virtual/VirtualDeviceManagerService.java
@@ -838,10 +838,11 @@
         }
 
         @Override
-        public boolean isDisplayOwnedByAnyVirtualDevice(int displayId) {
+        public boolean isInputDeviceOwnedByVirtualDevice(int inputDeviceId) {
             ArrayList<VirtualDeviceImpl> virtualDevicesSnapshot = getVirtualDevicesSnapshot();
             for (int i = 0; i < virtualDevicesSnapshot.size(); i++) {
-                if (virtualDevicesSnapshot.get(i).isDisplayOwnedByVirtualDevice(displayId)) {
+                if (virtualDevicesSnapshot.get(i)
+                        .isInputDeviceOwnedByVirtualDevice(inputDeviceId)) {
                     return true;
                 }
             }
diff --git a/services/core/java/com/android/server/companion/virtual/VirtualDeviceManagerInternal.java b/services/core/java/com/android/server/companion/virtual/VirtualDeviceManagerInternal.java
index 823788f..b179783 100644
--- a/services/core/java/com/android/server/companion/virtual/VirtualDeviceManagerInternal.java
+++ b/services/core/java/com/android/server/companion/virtual/VirtualDeviceManagerInternal.java
@@ -137,9 +137,9 @@
     public abstract boolean isAppRunningOnAnyVirtualDevice(int uid);
 
     /**
-     * Returns true if the {@code displayId} is owned by any virtual device
+     * @return whether the input device with the given id was created by a virtual device.
      */
-    public abstract boolean isDisplayOwnedByAnyVirtualDevice(int displayId);
+    public abstract boolean isInputDeviceOwnedByVirtualDevice(int inputDeviceId);
 
     /**
      * Gets the ids of VirtualDisplays owned by a VirtualDevice.
diff --git a/services/core/java/com/android/server/display/mode/VotesStatsReporter.java b/services/core/java/com/android/server/display/mode/VotesStatsReporter.java
index a30c4d2..e80b9451 100644
--- a/services/core/java/com/android/server/display/mode/VotesStatsReporter.java
+++ b/services/core/java/com/android/server/display/mode/VotesStatsReporter.java
@@ -16,12 +16,19 @@
 
 package com.android.server.display.mode;
 
+import static com.android.internal.util.FrameworkStatsLog.DISPLAY_MODE_DIRECTOR_VOTE_CHANGED;
+import static com.android.internal.util.FrameworkStatsLog.DISPLAY_MODE_DIRECTOR_VOTE_CHANGED__VOTE_STATUS__STATUS_ACTIVE;
+import static com.android.internal.util.FrameworkStatsLog.DISPLAY_MODE_DIRECTOR_VOTE_CHANGED__VOTE_STATUS__STATUS_ADDED;
+import static com.android.internal.util.FrameworkStatsLog.DISPLAY_MODE_DIRECTOR_VOTE_CHANGED__VOTE_STATUS__STATUS_REMOVED;
+
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.os.Trace;
 import android.util.SparseArray;
 import android.view.Display;
 
+import com.android.internal.util.FrameworkStatsLog;
+
 /**
  * The VotesStatsReporter is responsible for collecting and sending Vote related statistics 
  */
@@ -31,42 +38,77 @@
     private final boolean mIgnoredRenderRate;
     private final boolean mFrameworkStatsLogReportingEnabled;
 
+    private int mLastMinPriorityReported = Vote.MAX_PRIORITY + 1;
+
     public VotesStatsReporter(boolean ignoreRenderRate, boolean refreshRateVotingTelemetryEnabled) {
         mIgnoredRenderRate = ignoreRenderRate;
         mFrameworkStatsLogReportingEnabled = refreshRateVotingTelemetryEnabled;
     }
 
-    void reportVoteAdded(int displayId, int priority,  @NonNull Vote vote) {
+    void reportVoteChanged(int displayId, int priority,  @Nullable Vote vote) {
+        if (vote == null) {
+            reportVoteRemoved(displayId, priority);
+        } else {
+            reportVoteAdded(displayId, priority, vote);
+        }
+    }
+
+    private void reportVoteAdded(int displayId, int priority,  @NonNull Vote vote) {
         int maxRefreshRate = getMaxRefreshRate(vote, mIgnoredRenderRate);
         Trace.traceCounter(Trace.TRACE_TAG_POWER,
                 TAG + "." + displayId + ":" + Vote.priorityToString(priority), maxRefreshRate);
-        // if ( mFrameworkStatsLogReportingEnabled) {
-        // FrameworkStatsLog.write(VOTE_CHANGED, displayID, priority, ADDED, maxRefreshRate, -1);
-        // }
+        if (mFrameworkStatsLogReportingEnabled) {
+            FrameworkStatsLog.write(
+                    DISPLAY_MODE_DIRECTOR_VOTE_CHANGED, displayId, priority,
+                    DISPLAY_MODE_DIRECTOR_VOTE_CHANGED__VOTE_STATUS__STATUS_ADDED,
+                    maxRefreshRate, -1);
+        }
     }
 
-    void reportVoteRemoved(int displayId, int priority) {
+    private void reportVoteRemoved(int displayId, int priority) {
         Trace.traceCounter(Trace.TRACE_TAG_POWER,
                 TAG + "." + displayId + ":" + Vote.priorityToString(priority), -1);
-        // if ( mFrameworkStatsLogReportingEnabled) {
-        // FrameworkStatsLog.write(VOTE_CHANGED, displayID, priority, REMOVED, -1, -1);
-        // }
+        if (mFrameworkStatsLogReportingEnabled) {
+            FrameworkStatsLog.write(
+                    DISPLAY_MODE_DIRECTOR_VOTE_CHANGED, displayId, priority,
+                    DISPLAY_MODE_DIRECTOR_VOTE_CHANGED__VOTE_STATUS__STATUS_REMOVED, -1, -1);
+        }
     }
 
     void reportVotesActivated(int displayId, int minPriority, @Nullable Display.Mode baseMode,
             SparseArray<Vote> votes) {
-//        if (!mFrameworkStatsLogReportingEnabled) {
-//            return;
-//        }
-//        int selectedRefreshRate = baseMode != null ? (int) baseMode.getRefreshRate() : -1;
-//        for (int priority = minPriority; priority <= Vote.MAX_PRIORITY; priority ++) {
-//            Vote vote = votes.get(priority);
-//            if (vote != null) {
-//                int maxRefreshRate = getMaxRefreshRate(vote, mIgnoredRenderRate);
-//                FrameworkStatsLog.write(VOTE_CHANGED, displayId, priority,
-//                        ACTIVE, maxRefreshRate, selectedRefreshRate);
-//            }
-//        }
+        if (!mFrameworkStatsLogReportingEnabled) {
+            return;
+        }
+        int selectedRefreshRate = baseMode != null ? (int) baseMode.getRefreshRate() : -1;
+        for (int priority = Vote.MIN_PRIORITY; priority <= Vote.MAX_PRIORITY; priority++) {
+            if (priority < mLastMinPriorityReported && priority < minPriority) {
+                continue;
+            }
+            Vote vote = votes.get(priority);
+            if (vote == null) {
+                continue;
+            }
+
+            // Was previously reported ACTIVE, changed to ADDED
+            if (priority >= mLastMinPriorityReported && priority < minPriority) {
+                int maxRefreshRate = getMaxRefreshRate(vote, mIgnoredRenderRate);
+                FrameworkStatsLog.write(
+                        DISPLAY_MODE_DIRECTOR_VOTE_CHANGED, displayId, priority,
+                        DISPLAY_MODE_DIRECTOR_VOTE_CHANGED__VOTE_STATUS__STATUS_ADDED,
+                        maxRefreshRate, selectedRefreshRate);
+            }
+            // Was previously reported ADDED, changed to ACTIVE
+            if (priority >= minPriority && priority < mLastMinPriorityReported) {
+                int maxRefreshRate = getMaxRefreshRate(vote, mIgnoredRenderRate);
+                FrameworkStatsLog.write(
+                        DISPLAY_MODE_DIRECTOR_VOTE_CHANGED, displayId, priority,
+                        DISPLAY_MODE_DIRECTOR_VOTE_CHANGED__VOTE_STATUS__STATUS_ACTIVE,
+                        maxRefreshRate, selectedRefreshRate);
+            }
+
+            mLastMinPriorityReported = minPriority;
+        }
     }
 
     private static int getMaxRefreshRate(@NonNull Vote vote, boolean ignoreRenderRate) {
diff --git a/services/core/java/com/android/server/display/mode/VotesStorage.java b/services/core/java/com/android/server/display/mode/VotesStorage.java
index 7a1f7e9..56c7c18 100644
--- a/services/core/java/com/android/server/display/mode/VotesStorage.java
+++ b/services/core/java/com/android/server/display/mode/VotesStorage.java
@@ -117,22 +117,13 @@
             Slog.i(TAG, "Updated votes for display=" + displayId + " votes=" + votes);
         }
         if (changed) {
-            reportVoteStats(displayId, priority, vote);
+            if (mVotesStatsReporter != null) {
+                mVotesStatsReporter.reportVoteChanged(displayId, priority, vote);
+            }
             mListener.onChanged();
         }
     }
 
-    private void reportVoteStats(int displayId, int priority, @Nullable Vote vote) {
-        if (mVotesStatsReporter == null) {
-            return;
-        }
-        if (vote == null) {
-            mVotesStatsReporter.reportVoteRemoved(displayId, priority);
-        } else {
-            mVotesStatsReporter.reportVoteAdded(displayId, priority, vote);
-        }
-    }
-
     /** dump class values, for debugging */
     void dump(@NonNull PrintWriter pw) {
         SparseArray<SparseArray<Vote>> votesByDisplayLocal = new SparseArray<>();
diff --git a/services/core/java/com/android/server/webkit/WebViewUpdateServiceImpl2.java b/services/core/java/com/android/server/webkit/WebViewUpdateServiceImpl2.java
index 29782d9..f4fb1a1 100644
--- a/services/core/java/com/android/server/webkit/WebViewUpdateServiceImpl2.java
+++ b/services/core/java/com/android/server/webkit/WebViewUpdateServiceImpl2.java
@@ -159,11 +159,28 @@
         }
     }
 
+    private boolean shouldTriggerRepairLocked() {
+        if (mCurrentWebViewPackage == null) {
+            return true;
+        }
+        WebViewProviderInfo defaultProvider = getDefaultWebViewPackage();
+        if (mCurrentWebViewPackage.packageName.equals(defaultProvider.packageName)) {
+            List<UserPackage> userPackages =
+                    mSystemInterface.getPackageInfoForProviderAllUsers(
+                            mContext, defaultProvider);
+            return !isInstalledAndEnabledForAllUsers(userPackages);
+        } else {
+            return false;
+        }
+    }
+
     @Override
     public void prepareWebViewInSystemServer() {
         try {
+            boolean repairNeeded = true;
             synchronized (mLock) {
                 mCurrentWebViewPackage = findPreferredWebViewPackage();
+                repairNeeded = shouldTriggerRepairLocked();
                 String userSetting = mSystemInterface.getUserChosenWebViewProvider(mContext);
                 if (userSetting != null
                         && !userSetting.equals(mCurrentWebViewPackage.packageName)) {
@@ -177,26 +194,25 @@
                 }
                 onWebViewProviderChanged(mCurrentWebViewPackage);
             }
+
+            if (repairNeeded) {
+                // We didn't find a valid WebView implementation. Try explicitly re-enabling the
+                // default package for all users in case it was disabled, even if we already did the
+                // one-time migration before. If this actually changes the state, we will see the
+                // PackageManager broadcast shortly and try again.
+                WebViewProviderInfo defaultProvider = getDefaultWebViewPackage();
+                Slog.w(
+                        TAG,
+                        "No provider available for all users, trying to enable "
+                                + defaultProvider.packageName);
+                mSystemInterface.enablePackageForAllUsers(
+                        mContext, defaultProvider.packageName, true);
+            }
+
         } catch (Throwable t) {
             // Log and discard errors at this stage as we must not crash the system server.
             Slog.e(TAG, "error preparing webview provider from system server", t);
         }
-
-        if (getCurrentWebViewPackage() == null) {
-            // We didn't find a valid WebView implementation. Try explicitly re-enabling the
-            // fallback package for all users in case it was disabled, even if we already did the
-            // one-time migration before. If this actually changes the state, we will see the
-            // PackageManager broadcast shortly and try again.
-            WebViewProviderInfo[] webviewProviders = mSystemInterface.getWebViewPackages();
-            WebViewProviderInfo fallbackProvider = getFallbackProvider(webviewProviders);
-            if (fallbackProvider != null) {
-                Slog.w(TAG, "No valid provider, trying to enable " + fallbackProvider.packageName);
-                mSystemInterface.enablePackageForAllUsers(mContext, fallbackProvider.packageName,
-                                                          true);
-            } else {
-                Slog.e(TAG, "No valid provider and no fallback available.");
-            }
-        }
     }
 
     private void startZygoteWhenReady() {
@@ -421,42 +437,43 @@
 
     /**
      * Returns either the package info of the WebView provider determined in the following way:
-     * If the user has chosen a provider then use that if it is valid,
-     * otherwise use the first package in the webview priority list that is valid.
-     *
+     * If the user has chosen a provider then use that if it is valid, enabled and installed
+     * for all users, otherwise use the default provider.
      */
     private PackageInfo findPreferredWebViewPackage() throws WebViewPackageMissingException {
-        ProviderAndPackageInfo[] providers = getValidWebViewPackagesAndInfos();
-
-        String userChosenProvider = mSystemInterface.getUserChosenWebViewProvider(mContext);
-
         // If the user has chosen provider, use that (if it's installed and enabled for all
         // users).
-        for (ProviderAndPackageInfo providerAndPackage : providers) {
-            if (providerAndPackage.provider.packageName.equals(userChosenProvider)) {
-                // userPackages can contain null objects.
-                List<UserPackage> userPackages =
-                        mSystemInterface.getPackageInfoForProviderAllUsers(mContext,
-                                providerAndPackage.provider);
-                if (isInstalledAndEnabledForAllUsers(userPackages)) {
-                    return providerAndPackage.packageInfo;
+        String userChosenPackageName = mSystemInterface.getUserChosenWebViewProvider(mContext);
+        WebViewProviderInfo userChosenProvider =
+                getWebViewProviderForPackage(userChosenPackageName);
+        if (userChosenProvider != null) {
+            try {
+                PackageInfo packageInfo =
+                        mSystemInterface.getPackageInfoForProvider(userChosenProvider);
+                if (validityResult(userChosenProvider, packageInfo) == VALIDITY_OK) {
+                    List<UserPackage> userPackages =
+                            mSystemInterface.getPackageInfoForProviderAllUsers(
+                                    mContext, userChosenProvider);
+                    if (isInstalledAndEnabledForAllUsers(userPackages)) {
+                        return packageInfo;
+                    }
                 }
+            } catch (NameNotFoundException e) {
+                Slog.w(TAG, "User chosen WebView package (" + userChosenPackageName
+                        + ") not found");
             }
         }
 
-        // User did not choose, or the choice failed; use the most stable provider that is
-        // installed and enabled for all users, and available by default (not through
-        // user choice).
-        for (ProviderAndPackageInfo providerAndPackage : providers) {
-            if (providerAndPackage.provider.availableByDefault) {
-                // userPackages can contain null objects.
-                List<UserPackage> userPackages =
-                        mSystemInterface.getPackageInfoForProviderAllUsers(mContext,
-                                providerAndPackage.provider);
-                if (isInstalledAndEnabledForAllUsers(userPackages)) {
-                    return providerAndPackage.packageInfo;
-                }
+        // User did not choose, or the choice failed; return the default provider even if it is not
+        // installed or enabled for all users.
+        WebViewProviderInfo defaultProvider = getDefaultWebViewPackage();
+        try {
+            PackageInfo packageInfo = mSystemInterface.getPackageInfoForProvider(defaultProvider);
+            if (validityResult(defaultProvider, packageInfo) == VALIDITY_OK) {
+                return packageInfo;
             }
+        } catch (NameNotFoundException e) {
+            Slog.w(TAG, "Default WebView package (" + defaultProvider.packageName + ") not found");
         }
 
         // This should never happen during normal operation (only with modified system images).
@@ -464,6 +481,16 @@
         throw new WebViewPackageMissingException("Could not find a loadable WebView package");
     }
 
+    private WebViewProviderInfo getWebViewProviderForPackage(String packageName) {
+        WebViewProviderInfo[] allProviders = getWebViewPackages();
+        for (int n = 0; n < allProviders.length; n++) {
+            if (allProviders[n].packageName.equals(packageName)) {
+                return allProviders[n];
+            }
+        }
+        return null;
+    }
+
     /**
      * Return true iff {@param packageInfos} point to only installed and enabled packages.
      * The given packages {@param packageInfos} should all be pointing to the same package, but each
diff --git a/services/core/java/com/android/server/wm/BackNavigationController.java b/services/core/java/com/android/server/wm/BackNavigationController.java
index 8aaf76a..2bd49bf 100644
--- a/services/core/java/com/android/server/wm/BackNavigationController.java
+++ b/services/core/java/com/android/server/wm/BackNavigationController.java
@@ -1591,7 +1591,6 @@
 
     private static void setLaunchBehind(@NonNull ActivityRecord activity) {
         if (!activity.isVisibleRequested()) {
-            activity.setVisibility(true);
             // The transition could commit the visibility and in the finishing state, that could
             // skip commitVisibility call in setVisibility cause the activity won't visible here.
             // Call it again to make sure the activity could be visible while handling the pending
@@ -1669,10 +1668,14 @@
         ProtoLog.d(WM_DEBUG_BACK_PREVIEW, "onBackNavigationDone backType=%s, "
                 + "triggerBack=%b", backType, triggerBack);
 
-        mNavigationMonitor.stopMonitorForRemote();
-        mBackAnimationInProgress = false;
-        mShowWallpaper = false;
-        mPendingAnimationBuilder = null;
+        synchronized (mWindowManagerService.mGlobalLock) {
+            mNavigationMonitor.stopMonitorForRemote();
+            mBackAnimationInProgress = false;
+            mShowWallpaper = false;
+            // All animation should be done, clear any un-send animation.
+            mPendingAnimation = null;
+            mPendingAnimationBuilder = null;
+        }
     }
 
     static TaskSnapshot getSnapshot(@NonNull WindowContainer w,
diff --git a/services/credentials/java/com/android/server/credentials/CredentialManagerService.java b/services/credentials/java/com/android/server/credentials/CredentialManagerService.java
index fb0729f..667e086 100644
--- a/services/credentials/java/com/android/server/credentials/CredentialManagerService.java
+++ b/services/credentials/java/com/android/server/credentials/CredentialManagerService.java
@@ -63,7 +63,6 @@
 import android.util.Pair;
 import android.util.Slog;
 import android.util.SparseArray;
-import android.view.autofill.IAutoFillManagerClient;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.server.credentials.metrics.ApiName;
@@ -484,7 +483,7 @@
         public ICancellationSignal getCandidateCredentials(
                 GetCredentialRequest request,
                 IGetCandidateCredentialsCallback callback,
-                IAutoFillManagerClient clientCallback,
+                IBinder clientBinder,
                 final String callingPackage) {
             Slog.i(TAG, "starting getCandidateCredentials with callingPackage: "
                     + callingPackage);
@@ -506,7 +505,7 @@
                             constructCallingAppInfo(callingPackage, userId, request.getOrigin()),
                             getEnabledProvidersForUser(userId),
                             CancellationSignal.fromTransport(cancelTransport),
-                            clientCallback
+                            clientBinder
                     );
             addSessionLocked(userId, session);
 
diff --git a/services/credentials/java/com/android/server/credentials/GetCandidateRequestSession.java b/services/credentials/java/com/android/server/credentials/GetCandidateRequestSession.java
index 0187ce8..25281ba 100644
--- a/services/credentials/java/com/android/server/credentials/GetCandidateRequestSession.java
+++ b/services/credentials/java/com/android/server/credentials/GetCandidateRequestSession.java
@@ -30,11 +30,11 @@
 import android.credentials.ui.ProviderData;
 import android.credentials.ui.RequestInfo;
 import android.os.CancellationSignal;
+import android.os.IBinder;
 import android.os.RemoteException;
 import android.service.credentials.CallingAppInfo;
 import android.service.credentials.PermissionUtils;
 import android.util.Slog;
-import android.view.autofill.IAutoFillManagerClient;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -52,7 +52,7 @@
     private static final String SESSION_ID_KEY = "autofill_session_id";
     private static final String REQUEST_ID_KEY = "autofill_request_id";
 
-    private final IAutoFillManagerClient mAutoFillCallback;
+    private final IBinder mClientBinder;
     private final int mAutofillSessionId;
     private final int mAutofillRequestId;
 
@@ -62,15 +62,15 @@
             IGetCandidateCredentialsCallback callback, GetCredentialRequest request,
             CallingAppInfo callingAppInfo, Set<ComponentName> enabledProviders,
             CancellationSignal cancellationSignal,
-            IAutoFillManagerClient autoFillCallback) {
+            IBinder clientBinder) {
         super(context, sessionCallback, lock, userId, callingUid, request, callback,
                 RequestInfo.TYPE_GET, callingAppInfo, enabledProviders,
                 cancellationSignal, 0L, /*shouldBindClientToDeath=*/ false);
-        mAutoFillCallback = autoFillCallback;
+        mClientBinder = clientBinder;
         mAutofillSessionId = request.getData().getInt(SESSION_ID_KEY, -1);
         mAutofillRequestId = request.getData().getInt(REQUEST_ID_KEY, -1);
-        if (mAutoFillCallback != null) {
-            setUpClientCallbackListener(mAutoFillCallback.asBinder());
+        if (mClientBinder != null) {
+            setUpClientCallbackListener(mClientBinder);
         }
     }
 
diff --git a/services/foldables/devicestateprovider/proguard.flags b/services/foldables/devicestateprovider/proguard.flags
index 069cbc6..b810cad 100644
--- a/services/foldables/devicestateprovider/proguard.flags
+++ b/services/foldables/devicestateprovider/proguard.flags
@@ -1 +1 @@
--keep,allowoptimization,allowaccessmodification class com.android.server.policy.TentModeDeviceStatePolicy { *; }
+-keep,allowoptimization,allowaccessmodification class com.android.server.policy.BookStyleDeviceStatePolicy { *; }
diff --git a/services/foldables/devicestateprovider/src/com/android/server/policy/BookStyleClosedStatePredicate.java b/services/foldables/devicestateprovider/src/com/android/server/policy/BookStyleClosedStatePredicate.java
new file mode 100644
index 0000000..d5a3cff
--- /dev/null
+++ b/services/foldables/devicestateprovider/src/com/android/server/policy/BookStyleClosedStatePredicate.java
@@ -0,0 +1,432 @@
+/*
+ * 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.server.policy;
+
+import static android.hardware.SensorManager.SENSOR_DELAY_NORMAL;
+import static android.view.Display.DEFAULT_DISPLAY;
+
+import static com.android.server.policy.BookStylePreferredScreenCalculator.PreferredScreen.OUTER;
+import static com.android.server.policy.BookStylePreferredScreenCalculator.HingeAngle.ANGLE_0;
+import static com.android.server.policy.BookStylePreferredScreenCalculator.HingeAngle.ANGLE_0_TO_45;
+import static com.android.server.policy.BookStylePreferredScreenCalculator.HingeAngle.ANGLE_45_TO_90;
+import static com.android.server.policy.BookStylePreferredScreenCalculator.HingeAngle.ANGLE_90_TO_180;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.content.Context;
+import android.hardware.Sensor;
+import android.hardware.SensorEvent;
+import android.hardware.SensorEventListener;
+import android.hardware.SensorManager;
+import android.hardware.display.DisplayManager;
+import android.os.Handler;
+import android.util.ArraySet;
+import android.view.Display;
+import android.view.Surface;
+
+import com.android.server.policy.BookStylePreferredScreenCalculator.PreferredScreen;
+import com.android.server.policy.BookStylePreferredScreenCalculator.HingeAngle;
+import com.android.server.policy.BookStylePreferredScreenCalculator.StateTransition;
+import com.android.server.policy.BookStyleClosedStatePredicate.ConditionSensorListener.SensorSubscription;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
+import java.util.function.Predicate;
+import java.util.function.Supplier;
+
+/**
+ * 'Closed' state predicate that takes into account the posture of the device
+ * It accepts list of state transitions that control how the device moves between
+ * device states.
+ * See {@link BookStyleStateTransitions} for detailed description of the default behavior.
+ */
+public class BookStyleClosedStatePredicate implements Predicate<FoldableDeviceStateProvider>,
+        DisplayManager.DisplayListener {
+
+    private final BookStylePreferredScreenCalculator mClosedStateCalculator;
+    private final Handler mHandler = new Handler();
+    private final PostureEstimator mPostureEstimator;
+    private final DisplayManager mDisplayManager;
+
+    /**
+     * Creates {@link BookStyleClosedStatePredicate}. It is expected that the device has a pair
+     * of accelerometer sensors (one for each movable part of the device), see parameter
+     * descriptions for the behaviour when these sensors are not available.
+     * @param context context that could be used to get system services
+     * @param updatesListener callback that will be executed whenever the predicate should be
+     *                        checked again
+     * @param leftAccelerometerSensor accelerometer sensor that is located in the half of the
+     *                                device that has the outer screen, in case if this sensor is
+     *                                not provided, tent/wedge mode will be detected only using
+     *                                orientation sensor and screen rotation, so this mode won't
+     *                                be accessible by putting the device on a flat surface
+     * @param rightAccelerometerSensor accelerometer sensor that is located on the opposite side
+     *                                 across the hinge from the previous accelerometer sensor,
+     *                                 in case if this sensor is not provided, reverse wedge mode
+     *                                 won't be detected, so the device will use closed state using
+     *                                 constant angle when folding
+     * @param stateTransitions definition of all possible state transitions, see
+     *                         {@link BookStyleStateTransitions} for sample and more details
+     */
+
+    public BookStyleClosedStatePredicate(@NonNull Context context,
+            @NonNull ClosedStateUpdatesListener updatesListener,
+            @Nullable Sensor leftAccelerometerSensor, @Nullable Sensor rightAccelerometerSensor,
+            @NonNull List<StateTransition> stateTransitions) {
+        mDisplayManager = context.getSystemService(DisplayManager.class);
+        mDisplayManager.registerDisplayListener(this, mHandler);
+
+        mClosedStateCalculator = new BookStylePreferredScreenCalculator(stateTransitions);
+
+        final SensorManager sensorManager = context.getSystemService(SensorManager.class);
+        final Sensor orientationSensor = sensorManager.getDefaultSensor(
+                Sensor.TYPE_DEVICE_ORIENTATION);
+
+        mPostureEstimator = new PostureEstimator(mHandler, sensorManager,
+                leftAccelerometerSensor, rightAccelerometerSensor, orientationSensor,
+                updatesListener::onClosedStateUpdated);
+    }
+
+    /**
+     * Based on the current sensor readings and current state, returns true if the device should use
+     * 'CLOSED' device state and false if it should not use 'CLOSED' state (e.g. could use half-open
+     * or open states).
+     */
+    @Override
+    public boolean test(FoldableDeviceStateProvider foldableDeviceStateProvider) {
+        final HingeAngle hingeAngle = hingeAngleFromFloat(
+                foldableDeviceStateProvider.getHingeAngle());
+
+        mPostureEstimator.onDeviceClosedStatusChanged(hingeAngle == ANGLE_0);
+
+        final PreferredScreen preferredScreen = mClosedStateCalculator.
+                calculatePreferredScreen(hingeAngle, mPostureEstimator.isLikelyTentOrWedgeMode(),
+                        mPostureEstimator.isLikelyReverseWedgeMode(hingeAngle));
+
+        return preferredScreen == OUTER;
+    }
+
+    private HingeAngle hingeAngleFromFloat(float hingeAngle) {
+        if (hingeAngle == 0f) {
+            return ANGLE_0;
+        } else if (hingeAngle < 45f) {
+            return ANGLE_0_TO_45;
+        } else if (hingeAngle < 90f) {
+            return ANGLE_45_TO_90;
+        } else {
+            return ANGLE_90_TO_180;
+        }
+    }
+
+    @Override
+    public void onDisplayChanged(int displayId) {
+        if (displayId == DEFAULT_DISPLAY) {
+            final Display display = mDisplayManager.getDisplay(displayId);
+            int displayState = display.getState();
+            boolean isDisplayOn = displayState == Display.STATE_ON;
+            mPostureEstimator.onDisplayPowerStatusChanged(isDisplayOn);
+            mPostureEstimator.onDisplayRotationChanged(display.getRotation());
+        }
+    }
+
+    @Override
+    public void onDisplayAdded(int displayId) {
+
+    }
+
+    @Override
+    public void onDisplayRemoved(int displayId) {
+
+    }
+
+    public interface ClosedStateUpdatesListener {
+        void onClosedStateUpdated();
+    }
+
+    /**
+     * Estimates if the device is going to enter wedge/tent mode based on the sensor data
+     */
+    private static class PostureEstimator implements SensorEventListener {
+
+
+        private static final int FLAT_INCLINATION_THRESHOLD_DEGREES = 8;
+
+        /**
+         * Alpha parameter of the accelerometer low pass filter: the lower the value, the less high
+         * frequency noise it filter but reduces the latency.
+         */
+        private static final float GRAVITY_VECTOR_LOW_PASS_ALPHA_VALUE = 0.8f;
+
+
+        @Nullable
+        private final Sensor mLeftAccelerometerSensor;
+        @Nullable
+        private final Sensor mRightAccelerometerSensor;
+        private final Sensor mOrientationSensor;
+        private final Runnable mOnSensorUpdatedListener;
+
+        private final ConditionSensorListener mConditionedSensorListener;
+
+        @Nullable
+        private float[] mRightGravityVector;
+
+        @Nullable
+        private float[] mLeftGravityVector;
+
+        @Nullable
+        private Integer mLastScreenRotation;
+
+        @Nullable
+        private SensorEvent mLastDeviceOrientationSensorEvent = null;
+
+        private boolean mScreenTurnedOn = false;
+        private boolean mDeviceClosed = false;
+
+        public PostureEstimator(Handler handler, SensorManager sensorManager,
+                @Nullable Sensor leftAccelerometerSensor, @Nullable Sensor rightAccelerometerSensor,
+                Sensor orientationSensor, Runnable onSensorUpdated) {
+            mLeftAccelerometerSensor = leftAccelerometerSensor;
+            mRightAccelerometerSensor = rightAccelerometerSensor;
+            mOrientationSensor = orientationSensor;
+
+            mOnSensorUpdatedListener = onSensorUpdated;
+
+            final List<SensorSubscription> sensorSubscriptions = new ArrayList<>();
+            if (mLeftAccelerometerSensor != null) {
+                sensorSubscriptions.add(new SensorSubscription(
+                        mLeftAccelerometerSensor,
+                        /* allowedToListen= */ () -> mScreenTurnedOn && !mDeviceClosed,
+                        /* cleanup= */ () -> mLeftGravityVector = null));
+            }
+
+            if (mRightAccelerometerSensor != null) {
+                sensorSubscriptions.add(new SensorSubscription(
+                        mRightAccelerometerSensor,
+                        /* allowedToListen= */ () -> mScreenTurnedOn,
+                        /* cleanup= */ () -> mRightGravityVector = null));
+            }
+
+            sensorSubscriptions.add(new SensorSubscription(mOrientationSensor,
+                    /* allowedToListen= */ () -> mScreenTurnedOn,
+                    /* cleanup= */ () -> mLastDeviceOrientationSensorEvent = null));
+
+            mConditionedSensorListener = new ConditionSensorListener(sensorManager, this, handler,
+                    sensorSubscriptions);
+        }
+
+        @Override
+        public void onSensorChanged(SensorEvent event) {
+            if (event.sensor == mRightAccelerometerSensor) {
+                if (mRightGravityVector == null) {
+                    mRightGravityVector = new float[3];
+                }
+                setNewValueWithHighPassFilter(mRightGravityVector, event.values);
+
+                final boolean isRightMostlyFlat = Objects.equals(
+                        isGravityVectorMostlyFlat(mRightGravityVector), Boolean.TRUE);
+
+                if (isRightMostlyFlat) {
+                    // Reset orientation sensor when the device becomes flat
+                    mLastDeviceOrientationSensorEvent = null;
+                }
+            } else if (event.sensor == mLeftAccelerometerSensor) {
+                if (mLeftGravityVector == null) {
+                    mLeftGravityVector = new float[3];
+                }
+                setNewValueWithHighPassFilter(mLeftGravityVector, event.values);
+            } else if (event.sensor == mOrientationSensor) {
+                mLastDeviceOrientationSensorEvent = event;
+            }
+
+            mOnSensorUpdatedListener.run();
+        }
+
+        @Override
+        public void onAccuracyChanged(Sensor sensor, int accuracy) {
+
+        }
+
+        private void setNewValueWithHighPassFilter(float[] output, float[] newValues) {
+            final float alpha = GRAVITY_VECTOR_LOW_PASS_ALPHA_VALUE;
+            output[0] = alpha * output[0] + (1 - alpha) * newValues[0];
+            output[1] = alpha * output[1] + (1 - alpha) * newValues[1];
+            output[2] = alpha * output[2] + (1 - alpha) * newValues[2];
+        }
+
+        /**
+         * Returns true if the phone likely in reverse wedge mode (when a foldable phone is lying
+         * on the outer screen mostly flat to the ground)
+         */
+        public boolean isLikelyReverseWedgeMode(HingeAngle hingeAngle) {
+            return hingeAngle != ANGLE_0 && Objects.equals(
+                    isGravityVectorMostlyFlat(mLeftGravityVector), Boolean.TRUE);
+        }
+
+        /**
+         * Returns true if the phone is likely in tent or wedge mode when unfolding. Tent mode
+         * is detected by checking if the phone is in seascape position, screen is rotated to
+         * landscape or seascape, or if the right side of the device is mostly flat.
+         */
+        public boolean isLikelyTentOrWedgeMode() {
+            boolean isScreenLandscapeOrSeascape = Objects.equals(mLastScreenRotation,
+                    Surface.ROTATION_270) || Objects.equals(mLastScreenRotation,
+                    Surface.ROTATION_90);
+            if (isScreenLandscapeOrSeascape) {
+                return true;
+            }
+
+            boolean isRightMostlyFlat = Objects.equals(
+                    isGravityVectorMostlyFlat(mRightGravityVector), Boolean.TRUE);
+            if (isRightMostlyFlat) {
+                return true;
+            }
+
+            boolean isSensorSeaScape = Objects.equals(getOrientationSensorRotation(),
+                    Surface.ROTATION_270);
+            if (isSensorSeaScape) {
+                return true;
+            }
+
+            return false;
+        }
+
+        /**
+         * Returns true if the passed gravity vector implies that the phone is mostly flat (the
+         * vector is close to be perpendicular to the ground and has a positive Z component).
+         * Returns null if there is no data from the sensor.
+         */
+        private Boolean isGravityVectorMostlyFlat(@Nullable float[] vector) {
+            if (vector == null) return null;
+            if (vector[0] == 0.0f && vector[1] == 0.0f && vector[2] == 0.0f) {
+                // Likely we haven't received the actual data yet, treat it as no data
+                return null;
+            }
+
+            double vectorMagnitude = Math.sqrt(
+                    vector[0] * vector[0] + vector[1] * vector[1] + vector[2] * vector[2]);
+            float normalizedGravityZ = (float) (vector[2] / vectorMagnitude);
+
+            final int inclination = (int) Math.round(Math.toDegrees(Math.acos(normalizedGravityZ)));
+            return inclination < FLAT_INCLINATION_THRESHOLD_DEGREES;
+        }
+
+        private Integer getOrientationSensorRotation() {
+            if (mLastDeviceOrientationSensorEvent == null) return null;
+            return (int) mLastDeviceOrientationSensorEvent.values[0];
+        }
+
+        /**
+         * Called whenever display status changes, we use this signal to start/stop listening
+         * to sensors when the display is off to save battery. Using display state instead of
+         * general power state to reduce the time when sensors are on, we don't need to listen
+         * to the extra sensors when the screen is off.
+         */
+        public void onDisplayPowerStatusChanged(boolean screenTurnedOn) {
+            mScreenTurnedOn = screenTurnedOn;
+            mConditionedSensorListener.updateListeningState();
+        }
+
+        /**
+         * Called whenever we display rotation might have been updated
+         * @param rotation new rotation
+         */
+        public void onDisplayRotationChanged(int rotation) {
+            mLastScreenRotation = rotation;
+        }
+
+        /**
+         * Called whenever foldable device becomes fully closed or opened
+         */
+        public void onDeviceClosedStatusChanged(boolean deviceClosed) {
+            mDeviceClosed = deviceClosed;
+            mConditionedSensorListener.updateListeningState();
+        }
+    }
+
+    /**
+     * Helper class that subscribes or unsubscribes from a sensor based on a condition specified
+     * in {@link SensorSubscription}
+     */
+    static class ConditionSensorListener {
+        private final List<SensorSubscription> mSensorSubscriptions;
+        private final ArraySet<Sensor> mIsListening = new ArraySet<>();
+
+        private final SensorManager mSensorManager;
+        private final SensorEventListener mSensorEventListener;
+
+        private final Handler mHandler;
+
+        public ConditionSensorListener(SensorManager sensorManager,
+                SensorEventListener sensorEventListener, Handler handler,
+                List<SensorSubscription> sensorSubscriptions) {
+            mSensorManager = sensorManager;
+            mSensorEventListener = sensorEventListener;
+            mSensorSubscriptions = sensorSubscriptions;
+            mHandler = handler;
+        }
+
+        /**
+         * Updates current listening state of the sensor based on the provided conditions
+         */
+        public void updateListeningState() {
+            for (int i = 0; i < mSensorSubscriptions.size(); i++) {
+                final SensorSubscription subscription = mSensorSubscriptions.get(i);
+                final Sensor sensor = subscription.mSensor;
+
+                final boolean shouldBeListening = subscription.mAllowedToListenSupplier.get();
+                final boolean isListening = mIsListening.contains(sensor);
+                final boolean shouldUpdateListening = isListening != shouldBeListening;
+
+                if (shouldUpdateListening) {
+                    if (shouldBeListening) {
+                        mIsListening.add(sensor);
+                        mSensorManager.registerListener(mSensorEventListener, sensor,
+                                SENSOR_DELAY_NORMAL, mHandler);
+                    } else {
+                        mIsListening.remove(sensor);
+                        mSensorManager.unregisterListener(mSensorEventListener, sensor);
+                        subscription.mOnUnsubscribe.run();
+                    }
+                }
+            }
+        }
+
+        /**
+         * Represents a configuration of a single sensor subscription
+         */
+        public static class SensorSubscription {
+            private final Sensor mSensor;
+            private final Supplier<Boolean> mAllowedToListenSupplier;
+            private final Runnable mOnUnsubscribe;
+
+            /**
+             * @param sensor sensor to listen to
+             * @param allowedToListen return true when it is allowed to listen to the sensor
+             * @param cleanup a runnable that will be closed just before unsubscribing from the
+             *                sensor
+             */
+
+            public SensorSubscription(Sensor sensor, Supplier<Boolean> allowedToListen,
+                    Runnable cleanup) {
+                mSensor = sensor;
+                mAllowedToListenSupplier = allowedToListen;
+                mOnUnsubscribe = cleanup;
+            }
+        }
+    }
+}
diff --git a/services/foldables/devicestateprovider/src/com/android/server/policy/TentModeDeviceStatePolicy.java b/services/foldables/devicestateprovider/src/com/android/server/policy/BookStyleDeviceStatePolicy.java
similarity index 73%
rename from services/foldables/devicestateprovider/src/com/android/server/policy/TentModeDeviceStatePolicy.java
rename to services/foldables/devicestateprovider/src/com/android/server/policy/BookStyleDeviceStatePolicy.java
index 5968b63..ad938af 100644
--- a/services/foldables/devicestateprovider/src/com/android/server/policy/TentModeDeviceStatePolicy.java
+++ b/services/foldables/devicestateprovider/src/com/android/server/policy/BookStyleDeviceStatePolicy.java
@@ -21,10 +21,12 @@
 import static com.android.server.devicestate.DeviceState.FLAG_EMULATED_ONLY;
 import static com.android.server.devicestate.DeviceState.FLAG_UNSUPPORTED_WHEN_POWER_SAVE_MODE;
 import static com.android.server.devicestate.DeviceState.FLAG_UNSUPPORTED_WHEN_THERMAL_STATUS_CRITICAL;
+import static com.android.server.policy.BookStyleStateTransitions.DEFAULT_STATE_TRANSITIONS;
 import static com.android.server.policy.FoldableDeviceStateProvider.DeviceStateConfiguration.createConfig;
 import static com.android.server.policy.FoldableDeviceStateProvider.DeviceStateConfiguration.createTentModeClosedState;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.content.Context;
 import android.hardware.Sensor;
 import android.hardware.SensorManager;
@@ -39,12 +41,15 @@
 import java.util.function.Predicate;
 
 /**
- * Device state policy for a foldable device that supports tent mode: a mode when the device
- * keeps the outer display on until reaching a certain hinge angle threshold.
+ * Device state policy for a foldable device with two screens in a book style, where the hinge is
+ * located on the left side of the device when in folded posture.
+ * The policy supports tent/wedge mode: a mode when the device keeps the outer display on
+ * until reaching certain conditions like hinge angle threshold.
  *
  * Contains configuration for {@link FoldableDeviceStateProvider}.
  */
-public class TentModeDeviceStatePolicy extends DeviceStatePolicy {
+public class BookStyleDeviceStatePolicy extends DeviceStatePolicy implements
+        BookStyleClosedStatePredicate.ClosedStateUpdatesListener {
 
     private static final int DEVICE_STATE_CLOSED = 0;
     private static final int DEVICE_STATE_HALF_OPENED = 1;
@@ -57,9 +62,10 @@
     private static final int MIN_CLOSED_ANGLE_DEGREES = 0;
     private static final int MAX_CLOSED_ANGLE_DEGREES = 5;
 
-    private final DeviceStateProvider mProvider;
+    private final FoldableDeviceStateProvider mProvider;
 
     private final boolean mIsDualDisplayBlockingEnabled;
+    private final boolean mEnablePostureBasedClosedState;
     private static final Predicate<FoldableDeviceStateProvider> ALLOWED = p -> true;
     private static final Predicate<FoldableDeviceStateProvider> NOT_ALLOWED = p -> false;
 
@@ -73,30 +79,30 @@
      *                          between folded and unfolded modes, otherwise when folding the
      *                          display switch will happen at 0 degrees
      */
-    public TentModeDeviceStatePolicy(@NonNull Context context,
-            @NonNull Sensor hingeAngleSensor, @NonNull Sensor hallSensor, int closeAngleDegrees) {
-        this(new FeatureFlagsImpl(), context, hingeAngleSensor, hallSensor, closeAngleDegrees);
-    }
-
-    public TentModeDeviceStatePolicy(@NonNull FeatureFlags featureFlags, @NonNull Context context,
-                                     @NonNull Sensor hingeAngleSensor, @NonNull Sensor hallSensor,
-                                     int closeAngleDegrees) {
+    public BookStyleDeviceStatePolicy(@NonNull FeatureFlags featureFlags, @NonNull Context context,
+            @NonNull Sensor hingeAngleSensor, @NonNull Sensor hallSensor,
+            @Nullable Sensor leftAccelerometerSensor, @Nullable Sensor rightAccelerometerSensor,
+            Integer closeAngleDegrees) {
         super(context);
 
         final SensorManager sensorManager = mContext.getSystemService(SensorManager.class);
         final DisplayManager displayManager = mContext.getSystemService(DisplayManager.class);
 
-        final DeviceStateConfiguration[] configuration = createConfiguration(closeAngleDegrees);
-
+        mEnablePostureBasedClosedState = featureFlags.enableFoldablesPostureBasedClosedState();
         mIsDualDisplayBlockingEnabled = featureFlags.enableDualDisplayBlocking();
 
+        final DeviceStateConfiguration[] configuration = createConfiguration(
+                leftAccelerometerSensor, rightAccelerometerSensor, closeAngleDegrees);
+
         mProvider = new FoldableDeviceStateProvider(mContext, sensorManager,
                 hingeAngleSensor, hallSensor, displayManager, configuration);
     }
 
-    private DeviceStateConfiguration[] createConfiguration(int closeAngleDegrees) {
+    private DeviceStateConfiguration[] createConfiguration(@Nullable Sensor leftAccelerometerSensor,
+            @Nullable Sensor rightAccelerometerSensor, Integer closeAngleDegrees) {
         return new DeviceStateConfiguration[]{
-                createClosedConfiguration(closeAngleDegrees),
+                createClosedConfiguration(leftAccelerometerSensor, rightAccelerometerSensor,
+                        closeAngleDegrees),
                 createConfig(DEVICE_STATE_HALF_OPENED,
                         /* name= */ "HALF_OPENED",
                         /* activeStatePredicate= */ (provider) -> {
@@ -123,8 +129,10 @@
         };
     }
 
-    private DeviceStateConfiguration createClosedConfiguration(int closeAngleDegrees) {
-        if (closeAngleDegrees > 0) {
+    private DeviceStateConfiguration createClosedConfiguration(
+            @Nullable Sensor leftAccelerometerSensor, @Nullable Sensor rightAccelerometerSensor,
+            @Nullable Integer closeAngleDegrees) {
+        if (closeAngleDegrees != null) {
             // Switch displays at closeAngleDegrees in both ways (folding and unfolding)
             return createConfig(
                     DEVICE_STATE_CLOSED,
@@ -137,6 +145,19 @@
             );
         }
 
+        if (mEnablePostureBasedClosedState) {
+            // Use smart closed state predicate that will use different switch angles
+            // based on the device posture (e.g. wedge mode, tent mode, reverse wedge mode)
+            return createConfig(
+                    DEVICE_STATE_CLOSED,
+                    /* name= */ "CLOSED",
+                    /* flags= */ FLAG_CANCEL_OVERRIDE_REQUESTS,
+                    /* activeStatePredicate= */ new BookStyleClosedStatePredicate(mContext,
+                            this, leftAccelerometerSensor, rightAccelerometerSensor,
+                            DEFAULT_STATE_TRANSITIONS)
+            );
+        }
+
         // Switch to the outer display only at 0 degrees but use TENT_MODE_SWITCH_ANGLE_DEGREES
         // angle when switching to the inner display
         return createTentModeClosedState(DEVICE_STATE_CLOSED,
@@ -148,6 +169,11 @@
     }
 
     @Override
+    public void onClosedStateUpdated() {
+        mProvider.notifyDeviceStateChangedIfNeeded();
+    }
+
+    @Override
     public DeviceStateProvider getDeviceStateProvider() {
         return mProvider;
     }
diff --git a/services/foldables/devicestateprovider/src/com/android/server/policy/BookStylePreferredScreenCalculator.java b/services/foldables/devicestateprovider/src/com/android/server/policy/BookStylePreferredScreenCalculator.java
new file mode 100644
index 0000000..8977422
--- /dev/null
+++ b/services/foldables/devicestateprovider/src/com/android/server/policy/BookStylePreferredScreenCalculator.java
@@ -0,0 +1,309 @@
+/*
+ * 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.server.policy;
+
+import android.annotation.Nullable;
+
+import java.util.List;
+import java.util.Objects;
+
+/**
+ * Calculates if we should use outer or inner display on foldable devices based on a several
+ * inputs like device orientation, hinge angle signals.
+ *
+ * This is a stateful class and acts like a state machine with fixed number of states
+ * and transitions. It allows to list all possible state transitions instead of performing
+ * imperative logic to make sure that we cover all scenarios and improve debuggability.
+ *
+ * See {@link BookStyleStateTransitions} for detailed description of the default behavior.
+ */
+public class BookStylePreferredScreenCalculator {
+
+    /**
+     * When calculating the new state we will re-calculate it until it settles down. We re-calculate
+     * it because the new state might trigger another state transition and this might happen
+     * several times. We don't want to have infinite loops in state calculation, so this value
+     * limits the number of such state transitions.
+     * For example, in the default configuration {@link BookStyleStateTransitions}, after each
+     * transition with 'set sticky flag' output it will perform a transition to a state without
+     * 'set sticky flag' output.
+     * We also have a unit test covering all possible states which checks that we don't have such
+     * states that could end up in an infinite transition. See sample test for the default
+     * transitions in {@link BookStyleClosedStateCalculatorTest}.
+     */
+    private static final int MAX_STATE_CHANGES = 16;
+
+    private State mState = new State(
+            /* stickyKeepOuterUntil90Degrees= */ false,
+            /* stickyKeepInnerUntil45Degrees= */ false,
+            PreferredScreen.INVALID);
+
+    private final List<StateTransition> mStateTransitions;
+
+    /**
+     * Creates BookStyleClosedStateCalculator
+     * @param stateTransitions list of all state transitions
+     */
+    public BookStylePreferredScreenCalculator(List<StateTransition> stateTransitions) {
+        mStateTransitions = stateTransitions;
+    }
+
+    /**
+     * Calculates updated {@link PreferredScreen} based on the current inputs and the current state.
+     * The calculation is done based on defined {@link StateTransition}s, it might perform
+     * multiple transitions until we settle down on a single state. Multiple transitions could be
+     * performed in case if {@link StateTransition} causes another update of the state.
+     * There is a limit of maximum {@link MAX_STATE_CHANGES} state transitions, after which
+     * this method will throw an {@link IllegalStateException}.
+     *
+     * @param angle current hinge angle
+     * @param likelyTentOrWedge true if the device is likely in tent or wedge mode
+     * @param likelyReverseWedge true if the device is likely in reverse wedge mode
+     * @return updated {@link PreferredScreen}
+     */
+    public PreferredScreen calculatePreferredScreen(HingeAngle angle, boolean likelyTentOrWedge,
+            boolean likelyReverseWedge) {
+        int attempts = 0;
+        State newState = calculateNewState(mState, angle, likelyTentOrWedge, likelyReverseWedge);
+        while (attempts < MAX_STATE_CHANGES && !Objects.equals(mState, newState)) {
+            mState = newState;
+            newState = calculateNewState(mState, angle, likelyTentOrWedge, likelyReverseWedge);
+            attempts++;
+        }
+
+        if (attempts >= MAX_STATE_CHANGES) {
+            throw new IllegalStateException(
+                    "Can't settle state " + mState + ", inputs: hingeAngle = " + angle
+                            + ", likelyTentOrWedge = " + likelyTentOrWedge
+                            + ", likelyReverseWedge = " + likelyReverseWedge);
+        }
+
+        final State oldState = mState;
+        mState = newState;
+
+        if (mState.mPreferredScreen == PreferredScreen.INVALID) {
+            throw new IllegalStateException(
+                    "Reached invalid state " + mState + ", inputs: hingeAngle = " + angle
+                            + ", likelyTentOrWedge = " + likelyTentOrWedge
+                            + ", likelyReverseWedge = " + likelyReverseWedge + ", old state: "
+                            + oldState);
+        }
+
+        return mState.mPreferredScreen;
+    }
+
+    /**
+     * Returns the current state of the calculator
+     */
+    public State getState() {
+        return mState;
+    }
+
+    private State calculateNewState(State current, HingeAngle hingeAngle, boolean likelyTentOrWedge,
+            boolean likelyReverseWedge) {
+        for (int i = 0; i < mStateTransitions.size(); i++) {
+            final State newState = mStateTransitions.get(i).tryTransition(hingeAngle,
+                    likelyTentOrWedge, likelyReverseWedge, current);
+            if (newState != null) {
+                return newState;
+            }
+        }
+
+        throw new IllegalArgumentException(
+                "Entry not found for state: " + current + ", hingeAngle = " + hingeAngle
+                        + ", likelyTentOrWedge = " + likelyTentOrWedge + ", likelyReverseWedge = "
+                        + likelyReverseWedge);
+    }
+
+    /**
+     * The angle between two halves of the foldable device in degrees. The angle is '0' when
+     * the device is fully closed and '180' when the device is fully open and flat.
+     */
+    public enum HingeAngle {
+        ANGLE_0,
+        ANGLE_0_TO_45,
+        ANGLE_45_TO_90,
+        ANGLE_90_TO_180
+    }
+
+    /**
+     * Resulting closed state of the device, where OPEN state indicates that the device should use
+     * the inner display and CLOSED means that it should use the outer (cover) screen.
+     */
+    public enum PreferredScreen {
+        INNER,
+        OUTER,
+        INVALID
+    }
+
+    /**
+     * Describes a state transition for the posture based active screen calculator
+     */
+    public static class StateTransition {
+        private final Input mInput;
+        private final State mOutput;
+
+        public StateTransition(HingeAngle hingeAngle, boolean likelyTentOrWedge,
+                boolean likelyReverseWedge,
+                boolean stickyKeepOuterUntil90Degrees, boolean stickyKeepInnerUntil45Degrees,
+                PreferredScreen preferredScreen, Boolean setStickyKeepOuterUntil90Degrees,
+                Boolean setStickyKeepInnerUntil45Degrees) {
+            mInput = new Input(hingeAngle, likelyTentOrWedge, likelyReverseWedge,
+                    stickyKeepOuterUntil90Degrees, stickyKeepInnerUntil45Degrees);
+            mOutput = new State(setStickyKeepOuterUntil90Degrees,
+                    setStickyKeepInnerUntil45Degrees, preferredScreen);
+        }
+
+        /**
+         * Returns true if the state transition is applicable for the given inputs
+         */
+        private boolean isApplicable(HingeAngle hingeAngle, boolean likelyTentOrWedge,
+                boolean likelyReverseWedge, State currentState) {
+            return mInput.hingeAngle == hingeAngle
+                    && mInput.likelyTentOrWedge == likelyTentOrWedge
+                    && mInput.likelyReverseWedge == likelyReverseWedge
+                    && Objects.equals(mInput.stickyKeepOuterUntil90Degrees,
+                    currentState.stickyKeepOuterUntil90Degrees)
+                    && Objects.equals(mInput.stickyKeepInnerUntil45Degrees,
+                    currentState.stickyKeepInnerUntil45Degrees);
+        }
+
+        /**
+         * Try to perform transition for the inputs, returns new state if this
+         * transition is applicable for the given state and inputs
+         */
+        @Nullable
+        State tryTransition(HingeAngle hingeAngle, boolean likelyTentOrWedge,
+                boolean likelyReverseWedge, State currentState) {
+            if (!isApplicable(hingeAngle, likelyTentOrWedge, likelyReverseWedge, currentState)) {
+                return null;
+            }
+
+            boolean stickyKeepOuterUntil90Degrees = currentState.stickyKeepOuterUntil90Degrees;
+            boolean stickyKeepInnerUntil45Degrees = currentState.stickyKeepInnerUntil45Degrees;
+
+            if (mOutput.stickyKeepOuterUntil90Degrees != null) {
+                stickyKeepOuterUntil90Degrees =
+                        mOutput.stickyKeepOuterUntil90Degrees;
+            }
+
+            if (mOutput.stickyKeepInnerUntil45Degrees != null) {
+                stickyKeepInnerUntil45Degrees =
+                        mOutput.stickyKeepInnerUntil45Degrees;
+            }
+
+            return new State(stickyKeepOuterUntil90Degrees, stickyKeepInnerUntil45Degrees,
+                    mOutput.mPreferredScreen);
+        }
+    }
+
+    /**
+     * The input part of the {@link StateTransition}, these are the values that are used
+     * to decide which {@link State} output to choose.
+     */
+    private static class Input {
+        final HingeAngle hingeAngle;
+        final boolean likelyTentOrWedge;
+        final boolean likelyReverseWedge;
+        final boolean stickyKeepOuterUntil90Degrees;
+        final boolean stickyKeepInnerUntil45Degrees;
+
+        public Input(HingeAngle hingeAngle, boolean likelyTentOrWedge,
+                boolean likelyReverseWedge,
+                boolean stickyKeepOuterUntil90Degrees, boolean stickyKeepInnerUntil45Degrees) {
+            this.hingeAngle = hingeAngle;
+            this.likelyTentOrWedge = likelyTentOrWedge;
+            this.likelyReverseWedge = likelyReverseWedge;
+            this.stickyKeepOuterUntil90Degrees = stickyKeepOuterUntil90Degrees;
+            this.stickyKeepInnerUntil45Degrees = stickyKeepInnerUntil45Degrees;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (!(o instanceof Input)) return false;
+            Input that = (Input) o;
+            return likelyTentOrWedge == that.likelyTentOrWedge
+                    && likelyReverseWedge == that.likelyReverseWedge
+                    && stickyKeepOuterUntil90Degrees == that.stickyKeepOuterUntil90Degrees
+                    && stickyKeepInnerUntil45Degrees == that.stickyKeepInnerUntil45Degrees
+                    && hingeAngle == that.hingeAngle;
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(hingeAngle, likelyTentOrWedge, likelyReverseWedge,
+                    stickyKeepOuterUntil90Degrees, stickyKeepInnerUntil45Degrees);
+        }
+
+        @Override
+        public String toString() {
+            return "InputState{" +
+                    "hingeAngle=" + hingeAngle +
+                    ", likelyTentOrWedge=" + likelyTentOrWedge +
+                    ", likelyReverseWedge=" + likelyReverseWedge +
+                    ", stickyKeepOuterUntil90Degrees=" + stickyKeepOuterUntil90Degrees +
+                    ", stickyKeepInnerUntil45Degrees=" + stickyKeepInnerUntil45Degrees +
+                    '}';
+        }
+    }
+
+    /**
+     * Class that holds a state of the calculator, it could be used to store the current
+     * state or to define the target (output) state based on some input in {@link StateTransition}.
+     */
+    public static class State {
+        public Boolean stickyKeepOuterUntil90Degrees;
+        public Boolean stickyKeepInnerUntil45Degrees;
+
+        PreferredScreen mPreferredScreen;
+
+        public State(Boolean stickyKeepOuterUntil90Degrees,
+                Boolean stickyKeepInnerUntil45Degrees,
+                PreferredScreen preferredScreen) {
+            this.stickyKeepOuterUntil90Degrees = stickyKeepOuterUntil90Degrees;
+            this.stickyKeepInnerUntil45Degrees = stickyKeepInnerUntil45Degrees;
+            this.mPreferredScreen = preferredScreen;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (!(o instanceof State)) return false;
+            State that = (State) o;
+            return Objects.equals(stickyKeepOuterUntil90Degrees,
+                    that.stickyKeepOuterUntil90Degrees) && Objects.equals(
+                    stickyKeepInnerUntil45Degrees, that.stickyKeepInnerUntil45Degrees)
+                    && mPreferredScreen == that.mPreferredScreen;
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(stickyKeepOuterUntil90Degrees, stickyKeepInnerUntil45Degrees,
+                    mPreferredScreen);
+        }
+
+        @Override
+        public String toString() {
+            return "State{" +
+                    "stickyKeepOuterUntil90Degrees=" + stickyKeepOuterUntil90Degrees +
+                    ", stickyKeepInnerUntil90Degrees=" + stickyKeepInnerUntil45Degrees +
+                    ", closedState=" + mPreferredScreen +
+                    '}';
+        }
+    }
+}
diff --git a/services/foldables/devicestateprovider/src/com/android/server/policy/BookStyleStateTransitions.java b/services/foldables/devicestateprovider/src/com/android/server/policy/BookStyleStateTransitions.java
new file mode 100644
index 0000000..16daacb
--- /dev/null
+++ b/services/foldables/devicestateprovider/src/com/android/server/policy/BookStyleStateTransitions.java
@@ -0,0 +1,722 @@
+/*
+ * 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.server.policy;
+
+import com.android.server.policy.BookStylePreferredScreenCalculator.PreferredScreen;
+import com.android.server.policy.BookStylePreferredScreenCalculator.HingeAngle;
+import com.android.server.policy.BookStylePreferredScreenCalculator.StateTransition;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Describes all possible state transitions for {@link BookStylePreferredScreenCalculator}.
+ * It contains a default configuration for a foldable device that has two screens: smaller outer
+ * screen which has portrait natural orientation and a larger inner screen and allows to use the
+ * device in tent mode or wedge mode.
+ *
+ * As the output state could affect calculating of the new state, it could potentially cause
+ * infinite loop and make the state never settle down. This could be avoided using automated test
+ * that checks all possible inputs and asserts that the final state is valid.
+ * See sample test for the default transitions in {@link BookStyleClosedStateCalculatorTest}.
+ *
+ * - Tent mode is defined as a posture when the device is partially opened and placed on the ground
+ *   on the edges that are parallel to the hinge.
+ * - Wedge mode is when the device is partially opened and placed flat on the ground with the part
+ *   of the device that doesn't have the display
+ * - Reverse wedge mode is when the device is partially opened and placed flat on the ground with
+ *   the outer screen down, so the outer screen is not accessible
+ *
+ * Behavior description:
+ * - When unfolding with screens off we assume that no sensor data available except hinge angle
+ *   (based on hall sensor), so we switch to the inner screen immediately
+ *
+ * - When unfolding when screen is 'on' we can check if we are likely in tent or wedge mode
+ *   - If not likely tent/wedge mode or sensors data not available, then we unfold immediately
+ *     After unfolding, the state of the inner screen 'on' is sticky between 0 and 45 degrees, so
+ *     it won't jump back to the outer screen even if you move the phone into tent/wedge mode. The
+ *     stickiness is reset after fully closing the device or unfolding past 45 degrees.
+ *   - If likely tent or wedge mode, switch only at 90 degrees
+ *     Tent/wedge mode is 'sticky' between 0 and 90 degrees, so it won't reset until you either
+ *     fully close the device or unfold past 90 degrees.
+ *
+ * - When folding we can check if we are likely in reverse wedge mode
+ *   - If not likely in reverse wedge mode or sensor data is not available we switch to the outer
+ *     screen at 45 degrees and enable sticky tent/wedge mode as before, this allows to enter
+ *     tent/wedge mode even if you are not on an even surface or holding phone in landscape
+ *   - If likely in reverse wedge mode, switch to the outer screen only at 0 degrees to allow
+ *     some use cases like using camera in this posture, the check happens after passing 45 degrees
+ *     and inner screen becomes sticky turned 'on' until fully closing or unfolding past 45 degrees
+ */
+public class BookStyleStateTransitions {
+
+    public static final List<StateTransition> DEFAULT_STATE_TRANSITIONS = new ArrayList<>();
+
+    static {
+        // region Angle 0
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.OUTER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ true
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.OUTER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.OUTER,
+                /* setStickyKeepOuterUntil90Degrees */ false,
+                /* setStickyKeepInnerUntil45Degrees */ true
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INVALID,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.OUTER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ true
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.OUTER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.OUTER,
+                /* setStickyKeepOuterUntil90Degrees */ false,
+                /* setStickyKeepInnerUntil45Degrees */ true
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INVALID,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.OUTER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.OUTER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ false
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.OUTER,
+                /* setStickyKeepOuterUntil90Degrees */ false,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INVALID,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.OUTER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.OUTER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ false
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.OUTER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INVALID,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        // endregion
+
+        // region Angle 0-45
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0_TO_45,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.OUTER,
+                /* setStickyKeepOuterUntil90Degrees */ true,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0_TO_45,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0_TO_45,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.OUTER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0_TO_45,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INVALID,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0_TO_45,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ true
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0_TO_45,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0_TO_45,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.OUTER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0_TO_45,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INVALID,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0_TO_45,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.OUTER,
+                /* setStickyKeepOuterUntil90Degrees */ true,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0_TO_45,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0_TO_45,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.OUTER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0_TO_45,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INVALID,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0_TO_45,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ true
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0_TO_45,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0_TO_45,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.OUTER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_0_TO_45,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INVALID,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        // endregion
+
+        // region Angle 45-90
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_45_TO_90,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_45_TO_90,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ false
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_45_TO_90,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.OUTER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_45_TO_90,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INVALID,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_45_TO_90,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ true
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_45_TO_90,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_45_TO_90,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.OUTER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_45_TO_90,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INVALID,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_45_TO_90,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_45_TO_90,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ false
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_45_TO_90,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.OUTER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_45_TO_90,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INVALID,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_45_TO_90,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ true
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_45_TO_90,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_45_TO_90,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.OUTER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_45_TO_90,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INVALID,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        // endregion
+
+        // region Angle 90-180
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_90_TO_180,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_90_TO_180,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ false
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_90_TO_180,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ false,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_90_TO_180,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INVALID,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_90_TO_180,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_90_TO_180,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ false
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_90_TO_180,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ false,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_90_TO_180,
+                /* likelyTentOrWedge */ false,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INVALID,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_90_TO_180,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_90_TO_180,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ false
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_90_TO_180,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ false,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_90_TO_180,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ false,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INVALID,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_90_TO_180,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_90_TO_180,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ false,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ false
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_90_TO_180,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ false,
+                PreferredScreen.INNER,
+                /* setStickyKeepOuterUntil90Degrees */ false,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        DEFAULT_STATE_TRANSITIONS.add(new StateTransition(
+                HingeAngle.ANGLE_90_TO_180,
+                /* likelyTentOrWedge */ true,
+                /* likelyReverseWedge */ true,
+                /* stickyKeepOuterUntil90Degrees */ true,
+                /* stickyKeepInnerUntil45Degrees */ true,
+                PreferredScreen.INVALID,
+                /* setStickyKeepOuterUntil90Degrees */ null,
+                /* setStickyKeepInnerUntil45Degrees */ null
+        ));
+        // endregion
+    }
+}
diff --git a/services/foldables/devicestateprovider/tests/src/com/android/server/policy/BookStyleDeviceStatePolicyTest.java b/services/foldables/devicestateprovider/tests/src/com/android/server/policy/BookStyleDeviceStatePolicyTest.java
new file mode 100644
index 0000000..8d01b7a
--- /dev/null
+++ b/services/foldables/devicestateprovider/tests/src/com/android/server/policy/BookStyleDeviceStatePolicyTest.java
@@ -0,0 +1,703 @@
+/*
+ * 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.server.policy;
+
+import static android.view.Display.DEFAULT_DISPLAY;
+import static android.view.Display.STATE_OFF;
+import static android.view.Display.STATE_ON;
+
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.Truth.assertWithMessage;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.nullable;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.atLeastOnce;
+import static org.mockito.Mockito.clearInvocations;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import android.app.Instrumentation;
+import android.content.res.Configuration;
+import android.hardware.Sensor;
+import android.hardware.SensorEvent;
+import android.hardware.SensorEventListener;
+import android.hardware.SensorManager;
+import android.hardware.display.DisplayManager;
+import android.hardware.input.InputSensorInfo;
+import android.os.Handler;
+import android.testing.AndroidTestingRunner;
+import android.testing.TestableContext;
+import android.view.Display;
+import android.view.Surface;
+
+import androidx.test.platform.app.InstrumentationRegistry;
+
+import com.android.server.devicestate.DeviceStateProvider;
+import com.android.server.devicestate.DeviceStateProvider.Listener;
+import com.android.server.policy.feature.flags.FakeFeatureFlagsImpl;
+import com.android.server.policy.feature.flags.Flags;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Captor;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.mockito.internal.util.reflection.FieldSetter;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Unit tests for {@link BookStyleDeviceStatePolicy.Provider}.
+ * <p/>
+ * Run with <code>atest BookStyleDeviceStatePolicyTest</code>.
+ */
+@RunWith(AndroidTestingRunner.class)
+public final class BookStyleDeviceStatePolicyTest {
+
+    private static final int DEVICE_STATE_CLOSED = 0;
+    private static final int DEVICE_STATE_HALF_OPENED = 1;
+    private static final int DEVICE_STATE_OPENED = 2;
+
+    @Captor
+    private ArgumentCaptor<Integer> mDeviceStateCaptor;
+    @Captor
+    private ArgumentCaptor<DisplayManager.DisplayListener> mDisplayListenerCaptor;
+    @Mock
+    private SensorManager mSensorManager;
+    @Mock
+    private InputSensorInfo mInputSensorInfo;
+    @Mock
+    private Listener mListener;
+    @Mock
+    DisplayManager mDisplayManager;
+    @Mock
+    private Display mDisplay;
+
+    private final FakeFeatureFlagsImpl mFakeFeatureFlags = new FakeFeatureFlagsImpl();
+
+    private final Configuration mConfiguration = new Configuration();
+
+    private final Instrumentation mInstrumentation = InstrumentationRegistry.getInstrumentation();
+
+    @Rule
+    public final TestableContext mContext = new TestableContext(
+            mInstrumentation.getTargetContext());
+
+    private Sensor mHallSensor;
+    private Sensor mOrientationSensor;
+    private Sensor mHingeAngleSensor;
+    private Sensor mLeftAccelerometer;
+    private Sensor mRightAccelerometer;
+
+    private Map<Sensor, List<SensorEventListener>> mSensorEventListeners = new HashMap<>();
+    private DeviceStateProvider mProvider;
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+
+        mFakeFeatureFlags.setFlag(Flags.FLAG_ENABLE_FOLDABLES_POSTURE_BASED_CLOSED_STATE, true);
+        mFakeFeatureFlags.setFlag(Flags.FLAG_ENABLE_DUAL_DISPLAY_BLOCKING, true);
+
+        when(mInputSensorInfo.getName()).thenReturn("hall-effect");
+        mHallSensor = new Sensor(mInputSensorInfo);
+        when(mInputSensorInfo.getName()).thenReturn("hinge-angle");
+        mHingeAngleSensor = new Sensor(mInputSensorInfo);
+        when(mInputSensorInfo.getName()).thenReturn("left-accelerometer");
+        mLeftAccelerometer = new Sensor(mInputSensorInfo);
+        when(mInputSensorInfo.getName()).thenReturn("right-accelerometer");
+        mRightAccelerometer = new Sensor(mInputSensorInfo);
+        when(mInputSensorInfo.getName()).thenReturn("orientation");
+        mOrientationSensor = new Sensor(mInputSensorInfo);
+
+        mContext.addMockSystemService(SensorManager.class, mSensorManager);
+
+        when(mSensorManager.getDefaultSensor(eq(Sensor.TYPE_HINGE_ANGLE), eq(true)))
+                .thenReturn(mHingeAngleSensor);
+        when(mSensorManager.getDefaultSensor(eq(Sensor.TYPE_DEVICE_ORIENTATION)))
+                .thenReturn(mOrientationSensor);
+
+        when(mDisplayManager.getDisplay(eq(DEFAULT_DISPLAY))).thenReturn(mDisplay);
+        mContext.addMockSystemService(DisplayManager.class, mDisplayManager);
+
+        mContext.ensureTestableResources();
+        when(mContext.getResources().getConfiguration()).thenReturn(mConfiguration);
+
+        final List<Sensor> sensors = new ArrayList<>();
+        sensors.add(mHallSensor);
+        sensors.add(mHingeAngleSensor);
+        sensors.add(mOrientationSensor);
+        sensors.add(mLeftAccelerometer);
+        sensors.add(mRightAccelerometer);
+
+        when(mSensorManager.registerListener(any(), any(), anyInt(), any())).thenAnswer(
+                invocation -> {
+                    final SensorEventListener listener = invocation.getArgument(0);
+                    final Sensor sensor = invocation.getArgument(1);
+                    addSensorListener(sensor, listener);
+                    return true;
+                });
+        when(mSensorManager.registerListener(any(), any(), anyInt())).thenAnswer(
+                invocation -> {
+                    final SensorEventListener listener = invocation.getArgument(0);
+                    final Sensor sensor = invocation.getArgument(1);
+                    addSensorListener(sensor, listener);
+                    return true;
+                });
+
+        doAnswer(invocation -> {
+            final SensorEventListener listener = invocation.getArgument(0);
+            final boolean[] removed = {false};
+            mSensorEventListeners.forEach((sensor, sensorEventListeners) ->
+                    removed[0] |= sensorEventListeners.remove(listener));
+
+            if (!removed[0]) {
+                throw new IllegalArgumentException(
+                        "Trying to unregister listener " + listener + " that was not registered");
+            }
+
+            return null;
+        }).when(mSensorManager).unregisterListener(any(SensorEventListener.class));
+
+        doAnswer(invocation -> {
+            final SensorEventListener listener = invocation.getArgument(0);
+            final Sensor sensor = invocation.getArgument(1);
+
+            boolean removed = mSensorEventListeners.get(sensor).remove(listener);
+            if (!removed) {
+                throw new IllegalArgumentException(
+                        "Trying to unregister listener " + listener
+                                + " that was not registered for sensor " + sensor);
+            }
+
+            return null;
+        }).when(mSensorManager).unregisterListener(any(SensorEventListener.class),
+                any(Sensor.class));
+
+        try {
+            FieldSetter.setField(mHallSensor, mHallSensor.getClass()
+                    .getDeclaredField("mStringType"), "com.google.sensor.hall_effect");
+        } catch (NoSuchFieldException e) {
+            throw new RuntimeException(e);
+        }
+
+        when(mSensorManager.getSensorList(eq(Sensor.TYPE_ALL)))
+                .thenReturn(sensors);
+
+        mInstrumentation.runOnMainSync(() -> mProvider = createProvider());
+
+        verify(mDisplayManager, atLeastOnce()).registerDisplayListener(
+                mDisplayListenerCaptor.capture(), nullable(Handler.class));
+        setScreenOn(true);
+    }
+
+    @Test
+    public void test_noSensorEventsYet_reportOpenedState() {
+        mProvider.setListener(mListener);
+        verify(mListener).onStateChanged(mDeviceStateCaptor.capture());
+        assertEquals(DEVICE_STATE_OPENED, mDeviceStateCaptor.getValue().intValue());
+    }
+
+    @Test
+    public void test_deviceClosedSensorEventsBecameAvailable_reportsClosedState() {
+        mProvider.setListener(mListener);
+        clearInvocations(mListener);
+
+        sendHingeAngle(0f);
+
+        verify(mListener).onStateChanged(mDeviceStateCaptor.capture());
+        assertEquals(DEVICE_STATE_CLOSED, mDeviceStateCaptor.getValue().intValue());
+    }
+
+    @Test
+    public void test_hingeAngleClosed_reportsClosedState() {
+        sendHingeAngle(0f);
+
+        mProvider.setListener(mListener);
+        verify(mListener).onStateChanged(mDeviceStateCaptor.capture());
+        assertEquals(DEVICE_STATE_CLOSED, mDeviceStateCaptor.getValue().intValue());
+    }
+
+    @Test
+    public void test_hingeAngleFullyOpened_reportsOpenedState() {
+        sendHingeAngle(180f);
+
+        mProvider.setListener(mListener);
+        verify(mListener).onStateChanged(mDeviceStateCaptor.capture());
+        assertEquals(DEVICE_STATE_OPENED, mDeviceStateCaptor.getValue().intValue());
+    }
+
+    @Test
+    public void test_unfoldingFromClosedToFullyOpened_reportsOpenedEvent() {
+        sendHingeAngle(0f);
+        mProvider.setListener(mListener);
+        clearInvocations(mListener);
+
+        sendHingeAngle(180f);
+
+        verify(mListener).onStateChanged(mDeviceStateCaptor.capture());
+        assertEquals(DEVICE_STATE_OPENED, mDeviceStateCaptor.getValue().intValue());
+    }
+
+    @Test
+    public void test_foldingFromFullyOpenToFullyClosed_movesToClosedState() {
+        sendHingeAngle(180f);
+
+        sendHingeAngle(0f);
+
+        mProvider.setListener(mListener);
+        verify(mListener).onStateChanged(mDeviceStateCaptor.capture());
+        assertEquals(DEVICE_STATE_CLOSED, mDeviceStateCaptor.getValue().intValue());
+    }
+
+    @Test
+    public void test_slowUnfolding_reportsEventsInOrder() {
+        sendHingeAngle(0f);
+        mProvider.setListener(mListener);
+
+        sendHingeAngle(5f);
+        sendHingeAngle(10f);
+        sendHingeAngle(60f);
+        sendHingeAngle(100f);
+        sendHingeAngle(180f);
+
+        verify(mListener, atLeastOnce()).onStateChanged(mDeviceStateCaptor.capture());
+        assertThat(mDeviceStateCaptor.getAllValues()).containsExactly(
+                DEVICE_STATE_CLOSED,
+                DEVICE_STATE_HALF_OPENED,
+                DEVICE_STATE_OPENED
+        );
+    }
+
+    @Test
+    public void test_slowFolding_reportsEventsInOrder() {
+        sendHingeAngle(180f);
+        mProvider.setListener(mListener);
+
+        sendHingeAngle(180f);
+        sendHingeAngle(100f);
+        sendHingeAngle(60f);
+        sendHingeAngle(10f);
+        sendHingeAngle(5f);
+
+        verify(mListener, atLeastOnce()).onStateChanged(mDeviceStateCaptor.capture());
+        assertThat(mDeviceStateCaptor.getAllValues()).containsExactly(
+                DEVICE_STATE_OPENED,
+                DEVICE_STATE_HALF_OPENED,
+                DEVICE_STATE_CLOSED
+        );
+    }
+
+    @Test
+    public void test_hingeAngleOpen_screenOff_reportsHalfFolded() {
+        sendHingeAngle(0f);
+        setScreenOn(false);
+        mProvider.setListener(mListener);
+
+        sendHingeAngle(10f);
+
+        verify(mListener, atLeastOnce()).onStateChanged(mDeviceStateCaptor.capture());
+        assertThat(mDeviceStateCaptor.getAllValues()).containsExactly(
+                DEVICE_STATE_CLOSED,
+                DEVICE_STATE_HALF_OPENED
+        );
+    }
+
+    @Test
+    public void test_slowUnfoldingWithScreenOff_reportsEventsInOrder() {
+        sendHingeAngle(0f);
+        setScreenOn(false);
+        mProvider.setListener(mListener);
+
+        sendHingeAngle(5f);
+        assertLatestReportedState(DEVICE_STATE_HALF_OPENED);
+        sendHingeAngle(10f);
+        assertLatestReportedState(DEVICE_STATE_HALF_OPENED);
+        sendHingeAngle(60f);
+        assertLatestReportedState(DEVICE_STATE_HALF_OPENED);
+        sendHingeAngle(100f);
+        sendHingeAngle(180f);
+        assertLatestReportedState(DEVICE_STATE_OPENED);
+
+        verify(mListener, atLeastOnce()).onStateChanged(mDeviceStateCaptor.capture());
+        assertThat(mDeviceStateCaptor.getAllValues()).containsExactly(
+                DEVICE_STATE_CLOSED,
+                DEVICE_STATE_HALF_OPENED,
+                DEVICE_STATE_OPENED
+        );
+    }
+
+    @Test
+    public void test_unfoldWithScreenOff_reportsHalfOpened() {
+        sendHingeAngle(0f);
+        setScreenOn(false);
+        mProvider.setListener(mListener);
+
+        sendHingeAngle(5f);
+        sendHingeAngle(10f);
+
+        verify(mListener, atLeastOnce()).onStateChanged(mDeviceStateCaptor.capture());
+        assertThat(mDeviceStateCaptor.getAllValues()).containsExactly(
+                DEVICE_STATE_CLOSED,
+                DEVICE_STATE_HALF_OPENED
+        );
+    }
+
+    @Test
+    public void test_slowUnfoldingAndFolding_reportsEventsInOrder() {
+        sendHingeAngle(0f);
+        mProvider.setListener(mListener);
+        assertLatestReportedState(DEVICE_STATE_CLOSED);
+
+        // Started unfolding
+        sendHingeAngle(5f);
+        sendHingeAngle(30f);
+        assertLatestReportedState(DEVICE_STATE_HALF_OPENED);
+        sendHingeAngle(60f);
+        assertLatestReportedState(DEVICE_STATE_HALF_OPENED);
+        sendHingeAngle(100f);
+        assertLatestReportedState(DEVICE_STATE_HALF_OPENED);
+        sendHingeAngle(180f);
+        assertLatestReportedState(DEVICE_STATE_OPENED);
+
+        // Started folding
+        sendHingeAngle(100f);
+        assertLatestReportedState(DEVICE_STATE_HALF_OPENED);
+        sendHingeAngle(60f);
+        assertLatestReportedState(DEVICE_STATE_HALF_OPENED);
+        sendHingeAngle(30f);
+        assertLatestReportedState(DEVICE_STATE_CLOSED);
+        sendHingeAngle(5f);
+        assertLatestReportedState(DEVICE_STATE_CLOSED);
+
+        verify(mListener, atLeastOnce()).onStateChanged(mDeviceStateCaptor.capture());
+        assertThat(mDeviceStateCaptor.getAllValues()).containsExactly(
+                DEVICE_STATE_CLOSED,
+                DEVICE_STATE_HALF_OPENED,
+                DEVICE_STATE_OPENED,
+                DEVICE_STATE_HALF_OPENED,
+                DEVICE_STATE_CLOSED
+        );
+    }
+
+    @Test
+    public void test_unfoldTo30Degrees_screenOnRightSideMostlyFlat_keepsClosedState() {
+        sendHingeAngle(0f);
+        sendRightSideFlatSensorEvent(true);
+        mProvider.setListener(mListener);
+        assertLatestReportedState(DEVICE_STATE_CLOSED);
+        clearInvocations(mListener);
+
+        sendHingeAngle(30f);
+
+        verify(mListener, never()).onStateChanged(mDeviceStateCaptor.capture());
+    }
+
+    @Test
+    public void test_unfoldTo30Degrees_seascapeDeviceOrientation_keepsClosedState() {
+        sendHingeAngle(0f);
+        sendRightSideFlatSensorEvent(false);
+        sendDeviceOrientation(Surface.ROTATION_270);
+        mProvider.setListener(mListener);
+        assertLatestReportedState(DEVICE_STATE_CLOSED);
+        clearInvocations(mListener);
+
+        sendHingeAngle(30f);
+
+        verify(mListener, never()).onStateChanged(mDeviceStateCaptor.capture());
+    }
+
+    @Test
+    public void test_unfoldTo30Degrees_landscapeScreenRotation_keepsClosedState() {
+        sendHingeAngle(0f);
+        sendRightSideFlatSensorEvent(false);
+        sendScreenRotation(Surface.ROTATION_90);
+        mProvider.setListener(mListener);
+        assertLatestReportedState(DEVICE_STATE_CLOSED);
+        clearInvocations(mListener);
+
+        sendHingeAngle(30f);
+
+        verify(mListener, never()).onStateChanged(mDeviceStateCaptor.capture());
+    }
+
+    @Test
+    public void test_unfoldTo30Degrees_seascapeScreenRotation_keepsClosedState() {
+        sendHingeAngle(0f);
+        sendRightSideFlatSensorEvent(false);
+        sendScreenRotation(Surface.ROTATION_270);
+        mProvider.setListener(mListener);
+        assertLatestReportedState(DEVICE_STATE_CLOSED);
+        clearInvocations(mListener);
+
+        sendHingeAngle(30f);
+
+        verify(mListener, never()).onStateChanged(mDeviceStateCaptor.capture());
+    }
+
+    @Test
+    public void test_unfoldTo30Degrees_screenOnRightSideNotFlat_switchesToHalfOpenState() {
+        sendHingeAngle(0f);
+        sendRightSideFlatSensorEvent(false);
+        mProvider.setListener(mListener);
+        assertLatestReportedState(DEVICE_STATE_CLOSED);
+        clearInvocations(mListener);
+
+        sendHingeAngle(30f);
+
+        verify(mListener).onStateChanged(DEVICE_STATE_HALF_OPENED);
+    }
+
+    @Test
+    public void test_unfoldTo30Degrees_screenOffRightSideFlat_switchesToHalfOpenState() {
+        sendHingeAngle(0f);
+        setScreenOn(false);
+        // This sensor event should be ignored as screen is off
+        sendRightSideFlatSensorEvent(true);
+        mProvider.setListener(mListener);
+        assertLatestReportedState(DEVICE_STATE_CLOSED);
+        clearInvocations(mListener);
+
+        sendHingeAngle(30f);
+
+        verify(mListener).onStateChanged(DEVICE_STATE_HALF_OPENED);
+    }
+
+    @Test
+    public void test_unfoldTo60Degrees_andFoldTo10_switchesToClosedState() {
+        sendHingeAngle(0f);
+        sendRightSideFlatSensorEvent(false);
+        mProvider.setListener(mListener);
+        assertLatestReportedState(DEVICE_STATE_CLOSED);
+        sendHingeAngle(60f);
+        assertLatestReportedState(DEVICE_STATE_HALF_OPENED);
+        clearInvocations(mListener);
+
+        sendHingeAngle(10f);
+
+        verify(mListener).onStateChanged(DEVICE_STATE_CLOSED);
+    }
+
+    @Test
+    public void test_foldTo10AndUnfoldTo85Degrees_keepsClosedState() {
+        sendHingeAngle(0f);
+        sendRightSideFlatSensorEvent(false);
+        mProvider.setListener(mListener);
+        assertLatestReportedState(DEVICE_STATE_CLOSED);
+        sendHingeAngle(180f);
+        assertLatestReportedState(DEVICE_STATE_OPENED);
+        sendHingeAngle(10f);
+        assertLatestReportedState(DEVICE_STATE_CLOSED);
+
+        sendHingeAngle(85f);
+
+        // Keeps 'tent'/'wedge' mode even when right side is not flat
+        // as user manually folded the device not all the way
+        assertLatestReportedState(DEVICE_STATE_CLOSED);
+    }
+
+    @Test
+    public void test_foldTo0AndUnfoldTo85Degrees_doesNotKeepClosedState() {
+        sendHingeAngle(0f);
+        sendRightSideFlatSensorEvent(false);
+        mProvider.setListener(mListener);
+        assertLatestReportedState(DEVICE_STATE_CLOSED);
+        sendHingeAngle(180f);
+        assertLatestReportedState(DEVICE_STATE_OPENED);
+        sendHingeAngle(0f);
+        assertLatestReportedState(DEVICE_STATE_CLOSED);
+
+        sendHingeAngle(85f);
+
+        // Do not enter 'tent'/'wedge' mode when right side is not flat
+        // as user fully folded the device before that
+        assertLatestReportedState(DEVICE_STATE_HALF_OPENED);
+    }
+
+    @Test
+    public void test_foldTo10_leftSideIsFlat_keepsInnerScreenForReverseWedge() {
+        sendHingeAngle(180f);
+        sendLeftSideFlatSensorEvent(true);
+        mProvider.setListener(mListener);
+        assertLatestReportedState(DEVICE_STATE_OPENED);
+
+        sendHingeAngle(10f);
+
+        // Keep the inner screen for reverse wedge mode (e.g. for astrophotography use case)
+        assertLatestReportedState(DEVICE_STATE_HALF_OPENED);
+    }
+
+    @Test
+    public void test_foldTo10_leftSideIsNotFlat_switchesToOuterScreen() {
+        sendHingeAngle(180f);
+        sendLeftSideFlatSensorEvent(false);
+        mProvider.setListener(mListener);
+        assertLatestReportedState(DEVICE_STATE_OPENED);
+
+        sendHingeAngle(10f);
+
+        // Do not keep the inner screen as it is not reverse wedge mode
+        assertLatestReportedState(DEVICE_STATE_CLOSED);
+    }
+
+    @Test
+    public void test_foldTo10_noAccelerometerEvents_switchesToOuterScreen() {
+        sendHingeAngle(180f);
+        mProvider.setListener(mListener);
+        assertLatestReportedState(DEVICE_STATE_OPENED);
+
+        sendHingeAngle(10f);
+
+        // Do not keep the inner screen as it is not reverse wedge mode
+        assertLatestReportedState(DEVICE_STATE_CLOSED);
+    }
+
+    @Test
+    public void test_deviceClosed_screenIsOff_noSensorListeners() {
+        mProvider.setListener(mListener);
+
+        sendHingeAngle(0f);
+        setScreenOn(false);
+
+        assertNoListenersForSensor(mLeftAccelerometer);
+        assertNoListenersForSensor(mRightAccelerometer);
+        assertNoListenersForSensor(mOrientationSensor);
+    }
+
+    @Test
+    public void test_deviceClosed_screenIsOn_doesNotListenForOneAccelerometer() {
+        mProvider.setListener(mListener);
+
+        sendHingeAngle(0f);
+        setScreenOn(true);
+
+        assertNoListenersForSensor(mLeftAccelerometer);
+        assertListensForSensor(mRightAccelerometer);
+        assertListensForSensor(mOrientationSensor);
+    }
+
+    @Test
+    public void test_deviceOpened_screenIsOn_listensToSensors() {
+        mProvider.setListener(mListener);
+
+        sendHingeAngle(180f);
+        setScreenOn(true);
+
+        assertListensForSensor(mLeftAccelerometer);
+        assertListensForSensor(mRightAccelerometer);
+        assertListensForSensor(mOrientationSensor);
+    }
+
+    private void assertLatestReportedState(int state) {
+        final ArgumentCaptor<Integer> integerCaptor = ArgumentCaptor.forClass(Integer.class);
+        verify(mListener, atLeastOnce()).onStateChanged(integerCaptor.capture());
+        assertEquals(state, integerCaptor.getValue().intValue());
+    }
+
+    private void sendHingeAngle(float angle) {
+        sendSensorEvent(mHingeAngleSensor, new float[]{angle});
+    }
+
+    private void sendDeviceOrientation(int orientation) {
+        sendSensorEvent(mOrientationSensor, new float[]{orientation});
+    }
+
+    private void sendScreenRotation(int rotation) {
+        when(mDisplay.getRotation()).thenReturn(rotation);
+        mDisplayListenerCaptor.getAllValues().forEach((l) -> l.onDisplayChanged(DEFAULT_DISPLAY));
+    }
+
+    private void sendRightSideFlatSensorEvent(boolean flat) {
+        sendAccelerometerFlatEvents(mRightAccelerometer, flat);
+    }
+
+    private void sendLeftSideFlatSensorEvent(boolean flat) {
+        sendAccelerometerFlatEvents(mLeftAccelerometer, flat);
+    }
+
+    private static final int ACCELEROMETER_EVENTS = 10;
+
+    private void sendAccelerometerFlatEvents(Sensor sensor, boolean flat) {
+        final float[] values = flat ? new float[]{0.00021f, -0.00013f, 9.7899f} :
+                new float[]{6.124f, 4.411f, -1.7899f};
+        // Send the same values multiple times to bypass noise filter
+        for (int i = 0; i < ACCELEROMETER_EVENTS; i++) {
+            sendSensorEvent(sensor, values);
+        }
+    }
+
+    private void setScreenOn(boolean isOn) {
+        int state = isOn ? STATE_ON : STATE_OFF;
+        when(mDisplay.getState()).thenReturn(state);
+        mDisplayListenerCaptor.getAllValues().forEach((l) -> l.onDisplayChanged(DEFAULT_DISPLAY));
+    }
+
+    private void sendSensorEvent(Sensor sensor, float[] values) {
+        SensorEvent event = mock(SensorEvent.class);
+        event.sensor = sensor;
+        try {
+            FieldSetter.setField(event, event.getClass().getField("values"),
+                    values);
+        } catch (NoSuchFieldException e) {
+            throw new RuntimeException(e);
+        }
+
+        List<SensorEventListener> listeners = mSensorEventListeners.get(sensor);
+        if (listeners != null) {
+            listeners.forEach(sensorEventListener -> sensorEventListener.onSensorChanged(event));
+        }
+    }
+
+    private void assertNoListenersForSensor(Sensor sensor) {
+        final List<SensorEventListener> listeners = mSensorEventListeners.getOrDefault(sensor,
+                new ArrayList<>());
+        assertWithMessage("Expected no listeners for sensor " + sensor + " but found some").that(
+                listeners).isEmpty();
+    }
+
+    private void assertListensForSensor(Sensor sensor) {
+        final List<SensorEventListener> listeners = mSensorEventListeners.getOrDefault(sensor,
+                new ArrayList<>());
+        assertWithMessage(
+                "Expected at least one listener for sensor " + sensor).that(
+                listeners).isNotEmpty();
+    }
+
+    private void addSensorListener(Sensor sensor, SensorEventListener listener) {
+        List<SensorEventListener> listeners = mSensorEventListeners.computeIfAbsent(
+                sensor, k -> new ArrayList<>());
+        listeners.add(listener);
+    }
+
+    private DeviceStateProvider createProvider() {
+        return new BookStyleDeviceStatePolicy(mFakeFeatureFlags, mContext, mHingeAngleSensor,
+                mHallSensor, mLeftAccelerometer, mRightAccelerometer,
+                /* closeAngleDegrees= */ null).getDeviceStateProvider();
+    }
+}
diff --git a/services/foldables/devicestateprovider/tests/src/com/android/server/policy/BookStylePreferredScreenCalculatorTest.java b/services/foldables/devicestateprovider/tests/src/com/android/server/policy/BookStylePreferredScreenCalculatorTest.java
new file mode 100644
index 0000000..ae05b3f
--- /dev/null
+++ b/services/foldables/devicestateprovider/tests/src/com/android/server/policy/BookStylePreferredScreenCalculatorTest.java
@@ -0,0 +1,81 @@
+/*
+ * 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.server.policy;
+
+
+import static com.android.server.policy.BookStyleStateTransitions.DEFAULT_STATE_TRANSITIONS;
+
+import static com.google.common.truth.Truth.assertWithMessage;
+
+import android.testing.AndroidTestingRunner;
+
+import com.android.server.policy.BookStylePreferredScreenCalculator.PreferredScreen;
+import com.android.server.policy.BookStylePreferredScreenCalculator.HingeAngle;
+
+import com.google.common.collect.Lists;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.Arrays;
+import java.util.List;
+
+/**
+ * Unit tests for {@link BookStylePreferredScreenCalculator}.
+ * <p/>
+ * Run with <code>atest BookStyleClosedStateCalculatorTest</code>.
+ */
+@RunWith(AndroidTestingRunner.class)
+public final class BookStylePreferredScreenCalculatorTest {
+
+    private final BookStylePreferredScreenCalculator mCalculator =
+            new BookStylePreferredScreenCalculator(DEFAULT_STATE_TRANSITIONS);
+
+    private final List<HingeAngle> mHingeAngleValues = Arrays.asList(HingeAngle.values());
+    private final List<Boolean> mLikelyTentModeValues = Arrays.asList(true, false);
+    private final List<Boolean> mLikelyReverseWedgeModeValues = Arrays.asList(true, false);
+
+    @Test
+    public void transitionAllStates_noCrashes() {
+        final List<List<Object>> arguments = Lists.cartesianProduct(Arrays.asList(
+                mHingeAngleValues,
+                mLikelyTentModeValues,
+                mLikelyReverseWedgeModeValues
+        ));
+
+        arguments.forEach(objects -> {
+            final HingeAngle hingeAngle = (HingeAngle) objects.get(0);
+            final boolean likelyTent = (boolean) objects.get(1);
+            final boolean likelyReverseWedge = (boolean) objects.get(2);
+
+            final String description =
+                    "Input: hinge angle = " + hingeAngle + ", likelyTent = " + likelyTent
+                            + ", likelyReverseWedge = " + likelyReverseWedge;
+
+            // Verify that there are no crashes because of infinite state transitions and
+            // that it returns a valid active state
+            try {
+                PreferredScreen preferredScreen = mCalculator.calculatePreferredScreen(hingeAngle, likelyTent,
+                        likelyReverseWedge);
+
+                assertWithMessage(description).that(preferredScreen).isNotEqualTo(PreferredScreen.INVALID);
+            } catch (Throwable exception) {
+                throw new AssertionError(description, exception);
+            }
+        });
+    }
+}
diff --git a/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java b/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java
index e1f490a..5e7deef 100644
--- a/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java
@@ -608,6 +608,20 @@
     }
 
     @Test
+    public void testIsInputDeviceOwnedByVirtualDevice() {
+        assertThat(mLocalService.isInputDeviceOwnedByVirtualDevice(INPUT_DEVICE_ID)).isFalse();
+
+        final int fd = 1;
+        mInputController.addDeviceForTesting(BINDER, fd,
+                InputController.InputDeviceDescriptor.TYPE_KEYBOARD, DISPLAY_ID_1, PHYS,
+                DEVICE_NAME_1, INPUT_DEVICE_ID);
+        assertThat(mLocalService.isInputDeviceOwnedByVirtualDevice(INPUT_DEVICE_ID)).isTrue();
+
+        mInputController.unregisterInputDevice(BINDER);
+        assertThat(mLocalService.isInputDeviceOwnedByVirtualDevice(INPUT_DEVICE_ID)).isFalse();
+    }
+
+    @Test
     public void getDeviceIdsForUid_noRunningApps_returnsNull() {
         assertThat(mLocalService.getDeviceIdsForUid(UID_1)).isEmpty();
         assertThat(mVdmNative.getDeviceIdsForUid(UID_1)).isEmpty();
diff --git a/services/tests/servicestests/src/com/android/server/webkit/TestSystemImpl.java b/services/tests/servicestests/src/com/android/server/webkit/TestSystemImpl.java
index 3530e38..ae0a758 100644
--- a/services/tests/servicestests/src/com/android/server/webkit/TestSystemImpl.java
+++ b/services/tests/servicestests/src/com/android/server/webkit/TestSystemImpl.java
@@ -85,7 +85,7 @@
     private void enablePackageForUser(String packageName, boolean enable, int userId) {
         Map<Integer, PackageInfo> userPackages = mPackages.get(packageName);
         if (userPackages == null) {
-            throw new IllegalArgumentException("There is no package called " + packageName);
+            return;
         }
         PackageInfo packageInfo = userPackages.get(userId);
         packageInfo.applicationInfo.enabled = enable;
diff --git a/services/tests/servicestests/src/com/android/server/webkit/WebViewUpdateServiceTest.java b/services/tests/servicestests/src/com/android/server/webkit/WebViewUpdateServiceTest.java
index 32082e3..5a06327 100644
--- a/services/tests/servicestests/src/com/android/server/webkit/WebViewUpdateServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/webkit/WebViewUpdateServiceTest.java
@@ -127,12 +127,21 @@
     private void checkCertainPackageUsedAfterWebViewBootPreparation(String expectedProviderName,
             WebViewProviderInfo[] webviewPackages) {
         checkCertainPackageUsedAfterWebViewBootPreparation(
-                expectedProviderName, webviewPackages, 1);
+                expectedProviderName, webviewPackages, 1, null);
     }
 
     private void checkCertainPackageUsedAfterWebViewBootPreparation(String expectedProviderName,
-            WebViewProviderInfo[] webviewPackages, int numRelros) {
+            WebViewProviderInfo[] webviewPackages, String userSetting) {
+        checkCertainPackageUsedAfterWebViewBootPreparation(
+                expectedProviderName, webviewPackages, 1, userSetting);
+    }
+
+    private void checkCertainPackageUsedAfterWebViewBootPreparation(String expectedProviderName,
+            WebViewProviderInfo[] webviewPackages, int numRelros, String userSetting) {
         setupWithPackagesAndRelroCount(webviewPackages, numRelros);
+        if (userSetting != null) {
+            mTestSystemImpl.updateUserSetting(null, userSetting);
+        }
         // Add (enabled and valid) package infos for each provider
         setEnabledAndValidPackageInfos(webviewPackages);
 
@@ -280,7 +289,7 @@
                 singlePackage,
                 new WebViewProviderInfo[] {
                     new WebViewProviderInfo(singlePackage, "", true /*def av*/, false, null)},
-                2);
+                2, null);
     }
 
     // Ensure that package with valid signatures is chosen rather than package with invalid
@@ -295,14 +304,16 @@
         Signature invalidPackageSignature = new Signature("33");
 
         WebViewProviderInfo[] packages = new WebViewProviderInfo[] {
-            new WebViewProviderInfo(invalidPackage, "", true, false, new String[]{
-                        Base64.encodeToString(
-                                invalidExpectedSignature.toByteArray(), Base64.DEFAULT)}),
             new WebViewProviderInfo(validPackage, "", true, false, new String[]{
                         Base64.encodeToString(
-                                validSignature.toByteArray(), Base64.DEFAULT)})
+                                validSignature.toByteArray(), Base64.DEFAULT)}),
+            new WebViewProviderInfo(invalidPackage, "", true, false, new String[]{
+                        Base64.encodeToString(
+                                invalidExpectedSignature.toByteArray(), Base64.DEFAULT)})
         };
         setupWithPackagesNonDebuggable(packages);
+        // Start with the setting pointing to the invalid package
+        mTestSystemImpl.updateUserSetting(null, invalidPackage);
         mTestSystemImpl.setPackageInfo(createPackageInfo(invalidPackage, true /* enabled */,
                     true /* valid */, true /* installed */, new Signature[]{invalidPackageSignature}
                     , 0 /* updateTime */));
@@ -339,7 +350,9 @@
     }
 
     @Test
-    public void testFailListingEmptyWebviewPackages() {
+    @RequiresFlagsDisabled("android.webkit.update_service_v2")
+    // If the flag is set, will throw an exception because of no available by default provider.
+    public void testEmptyConfig() {
         WebViewProviderInfo[] packages = new WebViewProviderInfo[0];
         setupWithPackages(packages);
         setEnabledAndValidPackageInfos(packages);
@@ -352,14 +365,26 @@
         WebViewProviderResponse response = mWebViewUpdateServiceImpl.waitForAndGetProvider();
         assertEquals(WebViewFactory.LIBLOAD_FAILED_LISTING_WEBVIEW_PACKAGES, response.status);
         assertEquals(null, mWebViewUpdateServiceImpl.getCurrentWebViewPackage());
+    }
 
-        // Now install a package
+    @Test
+    public void testFailListingEmptyWebviewPackages() {
         String singlePackage = "singlePackage";
-        packages = new WebViewProviderInfo[]{
+        WebViewProviderInfo[] packages = new WebViewProviderInfo[]{
             new WebViewProviderInfo(singlePackage, "", true, false, null)};
         setupWithPackages(packages);
-        setEnabledAndValidPackageInfos(packages);
 
+        runWebViewBootPreparationOnMainSync();
+
+        Mockito.verify(mTestSystemImpl, Mockito.never()).onWebViewProviderChanged(
+                Matchers.anyObject());
+
+        WebViewProviderResponse response = mWebViewUpdateServiceImpl.waitForAndGetProvider();
+        assertEquals(WebViewFactory.LIBLOAD_FAILED_LISTING_WEBVIEW_PACKAGES, response.status);
+        assertEquals(null, mWebViewUpdateServiceImpl.getCurrentWebViewPackage());
+
+        // Now install the package
+        setEnabledAndValidPackageInfos(packages);
         mWebViewUpdateServiceImpl.packageStateChanged(singlePackage,
                 WebViewUpdateService.PACKAGE_ADDED, TestSystemImpl.PRIMARY_USER_ID);
 
@@ -370,7 +395,7 @@
         // Remove the package again
         mTestSystemImpl.removePackageInfo(singlePackage);
         mWebViewUpdateServiceImpl.packageStateChanged(singlePackage,
-                WebViewUpdateService.PACKAGE_ADDED, TestSystemImpl.PRIMARY_USER_ID);
+                WebViewUpdateService.PACKAGE_REMOVED, TestSystemImpl.PRIMARY_USER_ID);
 
         // Package removed - ensure our interface states that there is no package
         response = mWebViewUpdateServiceImpl.waitForAndGetProvider();
@@ -455,6 +480,8 @@
             new WebViewProviderInfo(firstPackage, "", true, false, null),
             new WebViewProviderInfo(secondPackage, "", true, false, null)};
         setupWithPackages(packages);
+        // Start with the setting pointing to the second package
+        mTestSystemImpl.updateUserSetting(null, secondPackage);
         // Have all packages be enabled, so that we can change provider however we want to
         setEnabledAndValidPackageInfos(packages);
 
@@ -463,9 +490,9 @@
         runWebViewBootPreparationOnMainSync();
 
         Mockito.verify(mTestSystemImpl).onWebViewProviderChanged(
-                Mockito.argThat(new IsPackageInfoWithName(firstPackage)));
+                Mockito.argThat(new IsPackageInfoWithName(secondPackage)));
 
-        assertEquals(firstPackage,
+        assertEquals(secondPackage,
                 mWebViewUpdateServiceImpl.getCurrentWebViewPackage().packageName);
 
         new Thread(new Runnable() {
@@ -474,12 +501,13 @@
                 WebViewProviderResponse threadResponse =
                     mWebViewUpdateServiceImpl.waitForAndGetProvider();
                 assertEquals(WebViewFactory.LIBLOAD_SUCCESS, threadResponse.status);
-                assertEquals(secondPackage, threadResponse.packageInfo.packageName);
-                // Verify that we killed the first package if we performed a settings change -
-                // otherwise we had to disable the first package, in which case its dependents
+                assertEquals(firstPackage, threadResponse.packageInfo.packageName);
+                // Verify that we killed the second package if we performed a settings change -
+                // otherwise we had to disable the second package, in which case its dependents
                 // should have been killed by the framework.
                 if (settingsChange) {
-                    Mockito.verify(mTestSystemImpl).killPackageDependents(Mockito.eq(firstPackage));
+                    Mockito.verify(mTestSystemImpl)
+                            .killPackageDependents(Mockito.eq(secondPackage));
                 }
                 countdown.countDown();
             }
@@ -490,32 +518,36 @@
         }
 
         if (settingsChange) {
-            mWebViewUpdateServiceImpl.changeProviderAndSetting(secondPackage);
+            mWebViewUpdateServiceImpl.changeProviderAndSetting(firstPackage);
         } else {
-            // Enable the second provider
-            mTestSystemImpl.setPackageInfo(createPackageInfo(secondPackage, true /* enabled */,
+            // Enable the first provider
+            mTestSystemImpl.setPackageInfo(createPackageInfo(firstPackage, true /* enabled */,
                         true /* valid */, true /* installed */));
             mWebViewUpdateServiceImpl.packageStateChanged(
-                    secondPackage, WebViewUpdateService.PACKAGE_CHANGED, TestSystemImpl.PRIMARY_USER_ID);
+                    firstPackage,
+                    WebViewUpdateService.PACKAGE_CHANGED,
+                    TestSystemImpl.PRIMARY_USER_ID);
 
             // Ensure we haven't changed package yet.
-            assertEquals(firstPackage,
+            assertEquals(secondPackage,
                     mWebViewUpdateServiceImpl.getCurrentWebViewPackage().packageName);
 
-            // Switch provider by disabling the first one
-            mTestSystemImpl.setPackageInfo(createPackageInfo(firstPackage, false /* enabled */,
+            // Switch provider by disabling the second one
+            mTestSystemImpl.setPackageInfo(createPackageInfo(secondPackage, false /* enabled */,
                         true /* valid */, true /* installed */));
             mWebViewUpdateServiceImpl.packageStateChanged(
-                    firstPackage, WebViewUpdateService.PACKAGE_CHANGED, TestSystemImpl.PRIMARY_USER_ID);
+                    secondPackage,
+                    WebViewUpdateService.PACKAGE_CHANGED,
+                    TestSystemImpl.PRIMARY_USER_ID);
         }
         mWebViewUpdateServiceImpl.notifyRelroCreationCompleted();
-        // first package done, should start on second
+        // second package done, should start on first
 
         Mockito.verify(mTestSystemImpl).onWebViewProviderChanged(
-                Mockito.argThat(new IsPackageInfoWithName(secondPackage)));
+                Mockito.argThat(new IsPackageInfoWithName(firstPackage)));
 
         mWebViewUpdateServiceImpl.notifyRelroCreationCompleted();
-        // second package done, the other thread should now be unblocked
+        // first package done, the other thread should now be unblocked
         try {
             countdown.await();
         } catch (InterruptedException e) {
@@ -526,6 +558,7 @@
      * Scenario for testing re-enabling a fallback package.
      */
     @Test
+    @RequiresFlagsDisabled("android.webkit.update_service_v2")
     public void testFallbackPackageEnabling() {
         String testPackage = "testFallback";
         WebViewProviderInfo[] packages = new WebViewProviderInfo[] {
@@ -555,6 +588,9 @@
      * 3. Primary should be used
      */
     @Test
+    @RequiresFlagsDisabled("android.webkit.update_service_v2")
+    // If the flag is set, we don't automitally switch to secondary package unless it is
+    // chosen directly.
     public void testInstallingPrimaryPackage() {
         String primaryPackage = "primary";
         String secondaryPackage = "secondary";
@@ -586,16 +622,16 @@
     }
 
     @Test
-    public void testRemovingPrimarySelectsSecondarySingleUser() {
+    public void testRemovingSecondarySelectsPrimarySingleUser() {
         for (PackageRemovalType removalType : REMOVAL_TYPES) {
-            checkRemovingPrimarySelectsSecondary(false /* multiUser */, removalType);
+            checkRemovingSecondarySelectsPrimary(false /* multiUser */, removalType);
         }
     }
 
     @Test
-    public void testRemovingPrimarySelectsSecondaryMultiUser() {
+    public void testRemovingSecondarySelectsPrimaryMultiUser() {
         for (PackageRemovalType removalType : REMOVAL_TYPES) {
-            checkRemovingPrimarySelectsSecondary(true /* multiUser */, removalType);
+            checkRemovingSecondarySelectsPrimary(true /* multiUser */, removalType);
         }
     }
 
@@ -609,7 +645,7 @@
 
     private PackageRemovalType[] REMOVAL_TYPES = PackageRemovalType.class.getEnumConstants();
 
-    public void checkRemovingPrimarySelectsSecondary(boolean multiUser,
+    private void checkRemovingSecondarySelectsPrimary(boolean multiUser,
             PackageRemovalType removalType) {
         String primaryPackage = "primary";
         String secondaryPackage = "secondary";
@@ -620,6 +656,8 @@
                     secondaryPackage, "", true /* default available */, false /* fallback */,
                     null)};
         setupWithPackages(packages);
+        // Start with the setting pointing to the secondary package
+        mTestSystemImpl.updateUserSetting(null, secondaryPackage);
         int secondaryUserId = 10;
         int userIdToChangePackageFor = multiUser ? secondaryUserId : TestSystemImpl.PRIMARY_USER_ID;
         if (multiUser) {
@@ -629,31 +667,31 @@
         setEnabledAndValidPackageInfosForUser(TestSystemImpl.PRIMARY_USER_ID, packages);
 
         runWebViewBootPreparationOnMainSync();
-        checkPreparationPhasesForPackage(primaryPackage, 1);
+        checkPreparationPhasesForPackage(secondaryPackage, 1);
 
         boolean enabled = !(removalType == PackageRemovalType.DISABLE);
         boolean installed = !(removalType == PackageRemovalType.UNINSTALL);
         boolean hidden = (removalType == PackageRemovalType.HIDE);
-        // Disable primary package and ensure secondary becomes used
+        // Disable secondary package and ensure primary becomes used
         mTestSystemImpl.setPackageInfoForUser(userIdToChangePackageFor,
-                createPackageInfo(primaryPackage, enabled /* enabled */, true /* valid */,
+                createPackageInfo(secondaryPackage, enabled /* enabled */, true /* valid */,
                     installed /* installed */, null /* signature */, 0 /* updateTime */,
                     hidden /* hidden */));
-        mWebViewUpdateServiceImpl.packageStateChanged(primaryPackage,
+        mWebViewUpdateServiceImpl.packageStateChanged(secondaryPackage,
                 removalType == PackageRemovalType.DISABLE
                 ? WebViewUpdateService.PACKAGE_CHANGED : WebViewUpdateService.PACKAGE_REMOVED,
                 userIdToChangePackageFor); // USER ID
-        checkPreparationPhasesForPackage(secondaryPackage, 1);
+        checkPreparationPhasesForPackage(primaryPackage, 1);
 
-        // Again enable primary package and verify primary is used
+        // Again enable secondary package and verify secondary is used
         mTestSystemImpl.setPackageInfoForUser(userIdToChangePackageFor,
-                createPackageInfo(primaryPackage, true /* enabled */, true /* valid */,
+                createPackageInfo(secondaryPackage, true /* enabled */, true /* valid */,
                     true /* installed */));
-        mWebViewUpdateServiceImpl.packageStateChanged(primaryPackage,
+        mWebViewUpdateServiceImpl.packageStateChanged(secondaryPackage,
                 removalType == PackageRemovalType.DISABLE
                 ? WebViewUpdateService.PACKAGE_CHANGED : WebViewUpdateService.PACKAGE_ADDED,
                 userIdToChangePackageFor);
-        checkPreparationPhasesForPackage(primaryPackage, 2);
+        checkPreparationPhasesForPackage(secondaryPackage, 2);
     }
 
     /**
@@ -671,18 +709,20 @@
                     secondaryPackage, "", true /* default available */, false /* fallback */,
                     null)};
         setupWithPackages(packages);
+        // Start with the setting pointing to the secondary package
+        mTestSystemImpl.updateUserSetting(null, secondaryPackage);
         setEnabledAndValidPackageInfosForUser(TestSystemImpl.PRIMARY_USER_ID, packages);
         int newUser = 100;
         mTestSystemImpl.addUser(newUser);
-        // Let the primary package be uninstalled for the new user
-        mTestSystemImpl.setPackageInfoForUser(newUser,
-                createPackageInfo(primaryPackage, true /* enabled */, true /* valid */,
-                        false /* installed */));
+        // Let the secondary package be uninstalled for the new user
         mTestSystemImpl.setPackageInfoForUser(newUser,
                 createPackageInfo(secondaryPackage, true /* enabled */, true /* valid */,
+                        false /* installed */));
+        mTestSystemImpl.setPackageInfoForUser(newUser,
+                createPackageInfo(primaryPackage, true /* enabled */, true /* valid */,
                         true /* installed */));
         mWebViewUpdateServiceImpl.handleNewUser(newUser);
-        checkPreparationPhasesForPackage(secondaryPackage, 1 /* numRelros */);
+        checkPreparationPhasesForPackage(primaryPackage, 1 /* numRelros */);
     }
 
     /**
@@ -780,9 +820,9 @@
         String chosenPackage = "chosenPackage";
         String nonChosenPackage = "non-chosenPackage";
         WebViewProviderInfo[] packages = new WebViewProviderInfo[] {
-            new WebViewProviderInfo(chosenPackage, "", true /* default available */,
-                    false /* fallback */, null),
             new WebViewProviderInfo(nonChosenPackage, "", true /* default available */,
+                    false /* fallback */, null),
+            new WebViewProviderInfo(chosenPackage, "", true /* default available */,
                     false /* fallback */, null)};
 
         setupWithPackages(packages);
@@ -810,6 +850,9 @@
     }
 
     @Test
+    @RequiresFlagsDisabled("android.webkit.update_service_v2")
+    // If the flag is set, we don't automitally switch to second package unless it is chosen
+    // directly.
     public void testRecoverFailedListingWebViewPackagesAddedPackage() {
         checkRecoverAfterFailListingWebviewPackages(false);
     }
@@ -874,22 +917,22 @@
                     false /* fallback */, null),
             new WebViewProviderInfo(secondPackage, "", true /* default available */,
                     false /* fallback */, null)};
-        checkCertainPackageUsedAfterWebViewBootPreparation(firstPackage, packages);
+        checkCertainPackageUsedAfterWebViewBootPreparation(secondPackage, packages, secondPackage);
 
         // Replace or remove the current webview package
         if (replaced) {
             mTestSystemImpl.setPackageInfo(
-                    createPackageInfo(firstPackage, true /* enabled */, false /* valid */,
+                    createPackageInfo(secondPackage, true /* enabled */, false /* valid */,
                         true /* installed */));
-            mWebViewUpdateServiceImpl.packageStateChanged(firstPackage,
+            mWebViewUpdateServiceImpl.packageStateChanged(secondPackage,
                     WebViewUpdateService.PACKAGE_ADDED_REPLACED, TestSystemImpl.PRIMARY_USER_ID);
         } else {
-            mTestSystemImpl.removePackageInfo(firstPackage);
-            mWebViewUpdateServiceImpl.packageStateChanged(firstPackage,
+            mTestSystemImpl.removePackageInfo(secondPackage);
+            mWebViewUpdateServiceImpl.packageStateChanged(secondPackage,
                     WebViewUpdateService.PACKAGE_REMOVED, TestSystemImpl.PRIMARY_USER_ID);
         }
 
-        checkPreparationPhasesForPackage(secondPackage, 1);
+        checkPreparationPhasesForPackage(firstPackage, 1);
 
         Mockito.verify(mTestSystemImpl, Mockito.never()).killPackageDependents(
                 Mockito.anyObject());
@@ -1073,10 +1116,12 @@
     }
 
     /**
-     * Ensure that the update service does use an uninstalled package when that is the only
+     * Ensure that the update service does not use an uninstalled package even if it is the only
      * package available.
      */
     @Test
+    @RequiresFlagsDisabled("android.webkit.update_service_v2")
+    // If the flag is set, we return the package even if it is not installed.
     public void testWithSingleUninstalledPackage() {
         String testPackageName = "test.package.name";
         WebViewProviderInfo[] webviewPackages = new WebViewProviderInfo[] {
@@ -1115,12 +1160,14 @@
         String installedPackage = "installedPackage";
         String uninstalledPackage = "uninstalledPackage";
         WebViewProviderInfo[] webviewPackages = new WebViewProviderInfo[] {
-            new WebViewProviderInfo(uninstalledPackage, "", true /* available by default */,
-                    false /* fallback */, null),
             new WebViewProviderInfo(installedPackage, "", true /* available by default */,
+                    false /* fallback */, null),
+            new WebViewProviderInfo(uninstalledPackage, "", true /* available by default */,
                     false /* fallback */, null)};
 
         setupWithPackages(webviewPackages);
+        // Start with the setting pointing to the uninstalled package
+        mTestSystemImpl.updateUserSetting(null, uninstalledPackage);
         int secondaryUserId = 5;
         if (multiUser) {
             mTestSystemImpl.addUser(secondaryUserId);
@@ -1128,7 +1175,7 @@
             setEnabledAndValidPackageInfosForUser(TestSystemImpl.PRIMARY_USER_ID, webviewPackages);
             mTestSystemImpl.setPackageInfoForUser(secondaryUserId, createPackageInfo(
                     installedPackage, true /* enabled */, true /* valid */, true /* installed */));
-            // Hide or uninstall the primary package for the second user
+            // Hide or uninstall the secondary package for the second user
             mTestSystemImpl.setPackageInfo(createPackageInfo(uninstalledPackage, true /* enabled */,
                     true /* valid */, (testUninstalled ? false : true) /* installed */,
                     null /* signatures */, 0 /* updateTime */, (testHidden ? true : false)));
@@ -1166,12 +1213,14 @@
         String installedPackage = "installedPackage";
         String uninstalledPackage = "uninstalledPackage";
         WebViewProviderInfo[] webviewPackages = new WebViewProviderInfo[] {
-            new WebViewProviderInfo(uninstalledPackage, "", true /* available by default */,
-                    false /* fallback */, null),
             new WebViewProviderInfo(installedPackage, "", true /* available by default */,
+                    false /* fallback */, null),
+            new WebViewProviderInfo(uninstalledPackage, "", true /* available by default */,
                     false /* fallback */, null)};
 
         setupWithPackages(webviewPackages);
+        // Start with the setting pointing to the uninstalled package
+        mTestSystemImpl.updateUserSetting(null, uninstalledPackage);
         int secondaryUserId = 412;
         mTestSystemImpl.addUser(secondaryUserId);
 
@@ -1221,12 +1270,14 @@
         String installedPackage = "installedPackage";
         String uninstalledPackage = "uninstalledPackage";
         WebViewProviderInfo[] webviewPackages = new WebViewProviderInfo[] {
-            new WebViewProviderInfo(uninstalledPackage, "", true /* available by default */,
-                    false /* fallback */, null),
             new WebViewProviderInfo(installedPackage, "", true /* available by default */,
+                    false /* fallback */, null),
+            new WebViewProviderInfo(uninstalledPackage, "", true /* available by default */,
                     false /* fallback */, null)};
 
         setupWithPackages(webviewPackages);
+        // Start with the setting pointing to the uninstalled package
+        mTestSystemImpl.updateUserSetting(null, uninstalledPackage);
         int secondaryUserId = 4;
         mTestSystemImpl.addUser(secondaryUserId);
 
@@ -1433,11 +1484,16 @@
                 new WebViewProviderInfo(newSdkPackage.packageName, "", true, false, null);
         WebViewProviderInfo currentSdkProviderInfo =
                 new WebViewProviderInfo(currentSdkPackage.packageName, "", true, false, null);
-        WebViewProviderInfo[] packages = new WebViewProviderInfo[] {
-            new WebViewProviderInfo(oldSdkPackage.packageName, "", true, false, null),
-            currentSdkProviderInfo, newSdkProviderInfo};
+        WebViewProviderInfo[] packages =
+                new WebViewProviderInfo[] {
+                    currentSdkProviderInfo,
+                    new WebViewProviderInfo(oldSdkPackage.packageName, "", true, false, null),
+                    newSdkProviderInfo
+                };
         setupWithPackages(packages);
-;
+        // Start with the setting pointing to the invalid package
+        mTestSystemImpl.updateUserSetting(null, oldSdkPackage.packageName);
+
         mTestSystemImpl.setPackageInfo(newSdkPackage);
         mTestSystemImpl.setPackageInfo(currentSdkPackage);
         mTestSystemImpl.setPackageInfo(oldSdkPackage);
@@ -1467,4 +1523,74 @@
         assertEquals(
                 defaultPackage1, mWebViewUpdateServiceImpl.getDefaultWebViewPackage().packageName);
     }
+
+    @Test
+    @RequiresFlagsEnabled("android.webkit.update_service_v2")
+    public void testDefaultWebViewPackageEnabling() {
+        String testPackage = "testDefault";
+        WebViewProviderInfo[] packages =
+                new WebViewProviderInfo[] {
+                    new WebViewProviderInfo(
+                            testPackage,
+                            "",
+                            true /* default available */,
+                            false /* fallback */,
+                            null)
+                };
+        setupWithPackages(packages);
+        mTestSystemImpl.setPackageInfo(
+                createPackageInfo(
+                        testPackage, false /* enabled */, true /* valid */, true /* installed */));
+
+        // Check that the boot time logic re-enables the default package.
+        runWebViewBootPreparationOnMainSync();
+        Mockito.verify(mTestSystemImpl)
+                .enablePackageForAllUsers(
+                        Matchers.anyObject(), Mockito.eq(testPackage), Mockito.eq(true));
+    }
+
+    private void testDefaultPackageChosen(PackageInfo packageInfo) {
+        WebViewProviderInfo[] packages =
+                new WebViewProviderInfo[] {
+                    new WebViewProviderInfo(packageInfo.packageName, "", true, false, null)
+                };
+        setupWithPackages(packages);
+        mTestSystemImpl.setPackageInfo(packageInfo);
+
+        runWebViewBootPreparationOnMainSync();
+        mWebViewUpdateServiceImpl.notifyRelroCreationCompleted();
+
+        assertEquals(
+                packageInfo.packageName,
+                mWebViewUpdateServiceImpl.getCurrentWebViewPackage().packageName);
+
+        WebViewProviderResponse response = mWebViewUpdateServiceImpl.waitForAndGetProvider();
+        assertEquals(packageInfo.packageName, response.packageInfo.packageName);
+    }
+
+    @Test
+    @RequiresFlagsEnabled("android.webkit.update_service_v2")
+    public void testDisabledDefaultPackageChosen() {
+        PackageInfo disabledPackage =
+                createPackageInfo(
+                        "disabledPackage",
+                        false /* enabled */,
+                        true /* valid */,
+                        true /* installed */);
+
+        testDefaultPackageChosen(disabledPackage);
+    }
+
+    @Test
+    @RequiresFlagsEnabled("android.webkit.update_service_v2")
+    public void testUninstalledDefaultPackageChosen() {
+        PackageInfo uninstalledPackage =
+                createPackageInfo(
+                        "uninstalledPackage",
+                        true /* enabled */,
+                        true /* valid */,
+                        false /* installed */);
+
+        testDefaultPackageChosen(uninstalledPackage);
+    }
 }
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
index 9bb2da0..ba7b52e 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
@@ -3365,7 +3365,7 @@
         assertFalse(app2.mActivityRecord.mImeInsetsFrozenUntilStartInput);
 
         if (Flags.bundleClientTransactionFlag()) {
-            verify(app2.getProcess()).scheduleClientTransactionItem(
+            verify(app2.getProcess(), atLeastOnce()).scheduleClientTransactionItem(
                     isA(WindowStateResizeItem.class));
         } else {
             verify(app2.mClient, atLeastOnce()).resized(any(), anyBoolean(), any(),