diff --git a/packages/SystemUI/aconfig/systemui.aconfig b/packages/SystemUI/aconfig/systemui.aconfig
index 3236130..d5278a7 100644
--- a/packages/SystemUI/aconfig/systemui.aconfig
+++ b/packages/SystemUI/aconfig/systemui.aconfig
@@ -357,3 +357,10 @@
    description: "Enables styled focus states on pin input field if keyboard is connected"
    bug: "316106516"
 }
+
+flag {
+    name: "keyguard_wm_state_refactor"
+    namespace: "systemui"
+    description: "Enables refactored logic for SysUI+WM unlock/occlusion code paths"
+    bug: "278086361"
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/keyguard/KeyguardSecurityContainerControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/keyguard/KeyguardSecurityContainerControllerTest.kt
index 2a793ea..0d2c64a 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/keyguard/KeyguardSecurityContainerControllerTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/keyguard/KeyguardSecurityContainerControllerTest.kt
@@ -202,11 +202,17 @@
         whenever(deviceProvisionedController.isUserSetup(anyInt())).thenReturn(true)
 
         featureFlags = FakeFeatureFlags()
-        featureFlags.set(Flags.KEYGUARD_WM_STATE_REFACTOR, false)
         featureFlags.set(Flags.REFACTOR_KEYGUARD_DISMISS_INTENT, false)
         featureFlags.set(Flags.LOCKSCREEN_ENABLE_LANDSCAPE, false)
 
-        mSetFlagsRule.enableFlags(AConfigFlags.FLAG_REVAMPED_BOUNCER_MESSAGES)
+        mSetFlagsRule.enableFlags(
+            AConfigFlags.FLAG_REVAMPED_BOUNCER_MESSAGES,
+        )
+        mSetFlagsRule.disableFlags(
+            FLAG_SIDEFPS_CONTROLLER_REFACTOR,
+            AConfigFlags.FLAG_KEYGUARD_WM_STATE_REFACTOR
+        )
+
         keyguardPasswordViewController =
             KeyguardPasswordViewController(
                 keyguardPasswordView,
@@ -237,7 +243,6 @@
         sceneInteractor.setTransitionState(sceneTransitionStateFlow)
         deviceEntryInteractor = kosmos.deviceEntryInteractor
 
-        mSetFlagsRule.disableFlags(FLAG_SIDEFPS_CONTROLLER_REFACTOR)
         underTest =
             KeyguardSecurityContainerController(
                 view,
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/data/repository/DeviceEntryFaceAuthRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/data/repository/DeviceEntryFaceAuthRepositoryTest.kt
index 6a14220..6808f5d 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/data/repository/DeviceEntryFaceAuthRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/data/repository/DeviceEntryFaceAuthRepositoryTest.kt
@@ -38,6 +38,7 @@
 import com.android.internal.logging.InstanceId.fakeInstanceId
 import com.android.internal.logging.UiEventLogger
 import com.android.keyguard.KeyguardUpdateMonitor
+import com.android.systemui.Flags as AConfigFlags
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.biometrics.data.repository.FakeDisplayStateRepository
 import com.android.systemui.biometrics.data.repository.FakeFacePropertyRepository
@@ -62,7 +63,6 @@
 import com.android.systemui.display.data.repository.display
 import com.android.systemui.dump.DumpManager
 import com.android.systemui.flags.FakeFeatureFlags
-import com.android.systemui.flags.Flags.KEYGUARD_WM_STATE_REFACTOR
 import com.android.systemui.keyguard.data.repository.BiometricType
 import com.android.systemui.keyguard.data.repository.FakeBiometricSettingsRepository
 import com.android.systemui.keyguard.data.repository.FakeCommandQueue
@@ -194,7 +194,7 @@
         biometricSettingsRepository = FakeBiometricSettingsRepository()
         deviceEntryFingerprintAuthRepository = FakeDeviceEntryFingerprintAuthRepository()
         trustRepository = FakeTrustRepository()
-        featureFlags = FakeFeatureFlags().apply { set(KEYGUARD_WM_STATE_REFACTOR, false) }
+        featureFlags = FakeFeatureFlags()
 
         powerRepository = FakePowerRepository()
         powerInteractor =
@@ -252,6 +252,10 @@
             .thenReturn(listOf(createFaceSensorProperties(supportsFaceDetection = true)))
         whenever(bypassController.bypassEnabled).thenReturn(true)
         underTest = createDeviceEntryFaceAuthRepositoryImpl(faceManager, bypassController)
+
+        mSetFlagsRule.disableFlags(
+            AConfigFlags.FLAG_KEYGUARD_WM_STATE_REFACTOR,
+        )
     }
 
     private fun createDeviceEntryFaceAuthRepositoryImpl(
@@ -301,7 +305,6 @@
             faceAuthBuffer,
             keyguardTransitionInteractor,
             displayStateInteractor,
-            featureFlags,
             dumpManager,
         )
     }
diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardSecurityContainerController.java b/packages/SystemUI/src/com/android/keyguard/KeyguardSecurityContainerController.java
index 8e5d0da..c23b3e9 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardSecurityContainerController.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardSecurityContainerController.java
@@ -19,9 +19,9 @@
 import static android.app.StatusBarManager.SESSION_KEYGUARD;
 import static android.content.res.Configuration.ORIENTATION_LANDSCAPE;
 
+import static com.android.keyguard.KeyguardSecurityContainer.BOUNCER_DISMISSIBLE_KEYGUARD;
 import static com.android.keyguard.KeyguardSecurityContainer.BOUNCER_DISMISS_BIOMETRIC;
 import static com.android.keyguard.KeyguardSecurityContainer.BOUNCER_DISMISS_EXTENDED_ACCESS;
-import static com.android.keyguard.KeyguardSecurityContainer.BOUNCER_DISMISSIBLE_KEYGUARD;
 import static com.android.keyguard.KeyguardSecurityContainer.BOUNCER_DISMISS_NONE_SECURITY;
 import static com.android.keyguard.KeyguardSecurityContainer.BOUNCER_DISMISS_PASSWORD;
 import static com.android.keyguard.KeyguardSecurityContainer.BOUNCER_DISMISS_SIM;
@@ -83,6 +83,7 @@
 import com.android.systemui.deviceentry.domain.interactor.DeviceEntryInteractor;
 import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.flags.Flags;
+import com.android.systemui.keyguard.KeyguardWmStateRefactor;
 import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor;
 import com.android.systemui.log.SessionTracker;
 import com.android.systemui.plugins.ActivityStarter;
@@ -99,8 +100,6 @@
 import com.android.systemui.util.kotlin.JavaAdapter;
 import com.android.systemui.util.settings.GlobalSettings;
 
-import dagger.Lazy;
-
 import java.io.File;
 import java.util.Arrays;
 import java.util.Optional;
@@ -108,6 +107,7 @@
 import javax.inject.Inject;
 import javax.inject.Provider;
 
+import dagger.Lazy;
 import kotlinx.coroutines.Job;
 
 /** Controller for {@link KeyguardSecurityContainer} */
@@ -329,7 +329,7 @@
                 }
             }
 
-            if (mFeatureFlags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+            if (KeyguardWmStateRefactor.isEnabled()) {
                 mKeyguardTransitionInteractor.startDismissKeyguardTransition();
             }
         }
diff --git a/packages/SystemUI/src/com/android/systemui/deviceentry/data/repository/DeviceEntryFaceAuthRepository.kt b/packages/SystemUI/src/com/android/systemui/deviceentry/data/repository/DeviceEntryFaceAuthRepository.kt
index 7a70c4a..cf7d601 100644
--- a/packages/SystemUI/src/com/android/systemui/deviceentry/data/repository/DeviceEntryFaceAuthRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/deviceentry/data/repository/DeviceEntryFaceAuthRepository.kt
@@ -40,8 +40,7 @@
 import com.android.systemui.deviceentry.shared.model.HelpFaceAuthenticationStatus
 import com.android.systemui.deviceentry.shared.model.SuccessFaceAuthenticationStatus
 import com.android.systemui.dump.DumpManager
-import com.android.systemui.flags.FeatureFlags
-import com.android.systemui.flags.Flags
+import com.android.systemui.keyguard.KeyguardWmStateRefactor
 import com.android.systemui.keyguard.data.repository.BiometricSettingsRepository
 import com.android.systemui.keyguard.data.repository.BiometricType
 import com.android.systemui.keyguard.data.repository.DeviceEntryFingerprintAuthRepository
@@ -63,10 +62,6 @@
 import com.android.systemui.user.data.model.SelectionStatus
 import com.android.systemui.user.data.repository.UserRepository
 import com.google.errorprone.annotations.CompileTimeConstant
-import java.io.PrintWriter
-import java.util.Arrays
-import java.util.stream.Collectors
-import javax.inject.Inject
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Job
@@ -88,6 +83,10 @@
 import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
+import java.io.PrintWriter
+import java.util.Arrays
+import java.util.stream.Collectors
+import javax.inject.Inject
 
 /**
  * API to run face authentication and detection for device entry / on keyguard (as opposed to the
@@ -165,7 +164,6 @@
     @FaceAuthTableLog private val faceAuthLog: TableLogBuffer,
     private val keyguardTransitionInteractor: KeyguardTransitionInteractor,
     private val displayStateInteractor: DisplayStateInteractor,
-    private val featureFlags: FeatureFlags,
     dumpManager: DumpManager,
 ) : DeviceEntryFaceAuthRepository, Dumpable {
     private var authCancellationSignal: CancellationSignal? = null
@@ -315,7 +313,7 @@
         // or device starts going to sleep.
         merge(
                 powerInteractor.isAsleep,
-                if (featureFlags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+                if (KeyguardWmStateRefactor.isEnabled) {
                     keyguardTransitionInteractor.isInTransitionToState(KeyguardState.GONE)
                 } else {
                     keyguardRepository.keyguardDoneAnimationsFinished.map { true }
diff --git a/packages/SystemUI/src/com/android/systemui/flags/Flags.kt b/packages/SystemUI/src/com/android/systemui/flags/Flags.kt
index ea8ba10..d797ec2 100644
--- a/packages/SystemUI/src/com/android/systemui/flags/Flags.kt
+++ b/packages/SystemUI/src/com/android/systemui/flags/Flags.kt
@@ -220,19 +220,6 @@
     @JvmField
     val WALLPAPER_PICKER_PREVIEW_ANIMATION = releasedFlag("wallpaper_picker_preview_animation")
 
-    /**
-     * TODO(b/278086361): Tracking bug
-     * Complete rewrite of the interactions between System UI and Window Manager involving keyguard
-     * state. When enabled, calls to ActivityTaskManagerService from System UI will exclusively
-     * occur from [WmLockscreenVisibilityManager] rather than the legacy KeyguardViewMediator.
-     *
-     * This flag is under development; some types of unlock may not animate properly if you enable
-     * it.
-     */
-    @JvmField
-    val KEYGUARD_WM_STATE_REFACTOR: UnreleasedFlag =
-            unreleasedFlag("keyguard_wm_state_refactor")
-
     // 300 - power menu
     // TODO(b/254512600): Tracking Bug
     @JvmField val POWER_MENU_LITE = releasedFlag("power_menu_lite")
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
index e2ab20e..f10b87e 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
@@ -77,7 +77,6 @@
 import com.android.systemui.SystemUIApplication;
 import com.android.systemui.dagger.qualifiers.Application;
 import com.android.systemui.flags.FeatureFlags;
-import com.android.systemui.flags.Flags;
 import com.android.systemui.keyguard.ui.binder.KeyguardSurfaceBehindParamsApplier;
 import com.android.systemui.keyguard.ui.binder.KeyguardSurfaceBehindViewBinder;
 import com.android.systemui.keyguard.ui.binder.WindowManagerLockscreenVisibilityViewBinder;
@@ -329,7 +328,7 @@
         mFlags = featureFlags;
         mPowerInteractor = powerInteractor;
 
-        if (mFlags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+        if (KeyguardWmStateRefactor.isEnabled()) {
             WindowManagerLockscreenVisibilityViewBinder.bind(
                     wmLockscreenVisibilityViewModel,
                     wmLockscreenVisibilityManager,
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardUnlockAnimationController.kt b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardUnlockAnimationController.kt
index 01ba0d2..53c81e5 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardUnlockAnimationController.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardUnlockAnimationController.kt
@@ -419,7 +419,7 @@
      */
     fun canPerformInWindowLauncherAnimations(): Boolean {
         // TODO(b/278086361): Refactor in-window animations.
-        return !featureFlags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR) &&
+        return !KeyguardWmStateRefactor.isEnabled &&
                 isSupportedLauncherUnderneath() &&
                 // If the launcher is underneath, but we're about to launch an activity, don't do
                 // the animations since they won't be visible.
@@ -866,7 +866,7 @@
         }
 
         surfaceBehindRemoteAnimationTargets?.forEach { surfaceBehindRemoteAnimationTarget ->
-            if (!featureFlags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+            if (!KeyguardWmStateRefactor.isEnabled) {
                 val surfaceHeight: Int =
                         surfaceBehindRemoteAnimationTarget.screenSpaceBounds.height()
 
@@ -1005,7 +1005,7 @@
         if (keyguardStateController.isShowing) {
             // Hide the keyguard, with no fade out since we animated it away during the unlock.
 
-            if (!featureFlags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+            if (!KeyguardWmStateRefactor.isEnabled) {
                 keyguardViewController.hide(
                         surfaceBehindRemoteAnimationStartTime,
                         0 /* fadeOutDuration */
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
index 5cebd96..691c1e2 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
@@ -139,7 +139,6 @@
 import com.android.systemui.dreams.DreamOverlayStateController;
 import com.android.systemui.dump.DumpManager;
 import com.android.systemui.flags.FeatureFlags;
-import com.android.systemui.flags.Flags;
 import com.android.systemui.flags.SystemPropertiesHelper;
 import com.android.systemui.keyguard.dagger.KeyguardModule;
 import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor;
@@ -175,8 +174,6 @@
 import com.android.systemui.wallpapers.data.repository.WallpaperRepository;
 import com.android.wm.shell.keyguard.KeyguardTransitions;
 
-import dagger.Lazy;
-
 import java.io.PrintWriter;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
@@ -186,6 +183,7 @@
 import java.util.concurrent.Executor;
 import java.util.function.Consumer;
 
+import dagger.Lazy;
 import kotlinx.coroutines.CoroutineDispatcher;
 
 /**
@@ -1051,7 +1049,7 @@
                 IRemoteAnimationFinishedCallback finishedCallback) {
             Trace.beginSection("mExitAnimationRunner.onAnimationStart#startKeyguardExitAnimation");
             startKeyguardExitAnimation(transit, apps, wallpapers, nonApps, finishedCallback);
-            if (mFeatureFlags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+            if (KeyguardWmStateRefactor.isEnabled()) {
                 mWmLockscreenVisibilityManager.get().onKeyguardGoingAwayRemoteAnimationStart(
                         transit, apps, wallpapers, nonApps, finishedCallback);
             }
@@ -1061,7 +1059,7 @@
         @Override // Binder interface
         public void onAnimationCancelled() {
             cancelKeyguardExitAnimation();
-            if (mFeatureFlags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+            if (KeyguardWmStateRefactor.isEnabled()) {
                 mWmLockscreenVisibilityManager.get().onKeyguardGoingAwayRemoteAnimationCancelled();
             }
         }
@@ -2737,7 +2735,7 @@
         mUiBgExecutor.execute(() -> {
             Log.d(TAG, "updateActivityLockScreenState(" + showing + ", " + aodShowing + ")");
 
-            if (mFeatureFlags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+            if (KeyguardWmStateRefactor.isEnabled()) {
                 // Handled in WmLockscreenVisibilityManager if flag is enabled.
                 return;
             }
@@ -2791,7 +2789,7 @@
             setShowingLocked(true, hidingOrGoingAway /* force */);
             mHiding = false;
 
-            if (!mFeatureFlags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+            if (!KeyguardWmStateRefactor.isEnabled()) {
                 // Handled directly in StatusBarKeyguardViewManager if enabled.
                 mKeyguardViewControllerLazy.get().show(options);
             }
@@ -2868,7 +2866,7 @@
             mKeyguardViewControllerLazy.get().setKeyguardGoingAwayState(true);
 
             // Handled in WmLockscreenVisibilityManager if flag is enabled.
-            if (!mFeatureFlags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+            if (!KeyguardWmStateRefactor.isEnabled()) {
                     // Don't actually hide the Keyguard at the moment, wait for window manager 
                     // until it tells us it's safe to do so with startKeyguardExitAnimation.
 		    // Posting to mUiOffloadThread to ensure that calls to ActivityTaskManager 
@@ -2974,7 +2972,7 @@
             } else {
                 Log.d(TAG, "Hiding keyguard while occluded. Just hide the keyguard view and exit.");
 
-                if (!mFeatureFlags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+                if (!KeyguardWmStateRefactor.isEnabled()) {
                     mKeyguardViewControllerLazy.get().hide(
                             mSystemClock.uptimeMillis() + mHideAnimation.getStartOffset(),
                             mHideAnimation.getDuration());
@@ -3010,7 +3008,7 @@
                 // If the flag is enabled, remote animation state is handled in
                 // WmLockscreenVisibilityManager.
                 if (finishedCallback != null
-                        && !mFeatureFlags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+                        && !KeyguardWmStateRefactor.isEnabled()) {
                     // There will not execute animation, send a finish callback to ensure the remote
                     // animation won't hang there.
                     try {
@@ -3036,7 +3034,7 @@
                         new IRemoteAnimationFinishedCallback() {
                             @Override
                             public void onAnimationFinished() throws RemoteException {
-                                if (!mFeatureFlags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+                                if (!KeyguardWmStateRefactor.isEnabled()) {
                                     try {
                                         finishedCallback.onAnimationFinished();
                                     } catch (RemoteException e) {
@@ -3068,7 +3066,7 @@
             // it will dismiss the panel in that case.
             } else if (!mStatusBarStateController.leaveOpenOnKeyguardHide()
                     && apps != null && apps.length > 0) {
-                if (!mFeatureFlags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+                if (!KeyguardWmStateRefactor.isEnabled()) {
                     // Handled in WmLockscreenVisibilityManager. Other logic in this class will
                     // short circuit when this is null.
                     mSurfaceBehindRemoteAnimationFinishedCallback = finishedCallback;
@@ -3092,7 +3090,7 @@
                         createInteractionJankMonitorConf(
                                 CUJ_LOCKSCREEN_UNLOCK_ANIMATION, "RemoteAnimationDisabled"));
 
-                if (!mFeatureFlags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+                if (!KeyguardWmStateRefactor.isEnabled()) {
                     // Handled directly in StatusBarKeyguardViewManager if enabled.
                     mKeyguardViewControllerLazy.get().hide(startTime, fadeoutDuration);
                 }
@@ -3111,7 +3109,7 @@
                         Slog.e(TAG, "Keyguard exit without a corresponding app to show.");
 
                         try {
-                            if (!mFeatureFlags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+                            if (!KeyguardWmStateRefactor.isEnabled()) {
                                 finishedCallback.onAnimationFinished();
                             }
                         } catch (RemoteException e) {
@@ -3143,7 +3141,7 @@
                         @Override
                         public void onAnimationEnd(Animator animation) {
                             try {
-                                if (!mFeatureFlags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+                                if (!KeyguardWmStateRefactor.isEnabled()) {
                                     finishedCallback.onAnimationFinished();
                                 }
                             } catch (RemoteException e) {
@@ -3156,7 +3154,7 @@
                         @Override
                         public void onAnimationCancel(Animator animation) {
                             try {
-                                if (!mFeatureFlags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+                                if (!KeyguardWmStateRefactor.isEnabled()) {
                                     finishedCallback.onAnimationFinished();
                                 }
                             } catch (RemoteException e) {
@@ -3321,7 +3319,7 @@
                 flags |= KEYGUARD_GOING_AWAY_FLAG_TO_LAUNCHER_CLEAR_SNAPSHOT;
             }
 
-            if (!mFeatureFlags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+            if (!KeyguardWmStateRefactor.isEnabled()) {
                 // Handled in WmLockscreenVisibilityManager.
                 mActivityTaskManagerService.keyguardGoingAway(flags);
             }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardWmStateRefactor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardWmStateRefactor.kt
new file mode 100644
index 0000000..ddccc5d
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardWmStateRefactor.kt
@@ -0,0 +1,53 @@
+/*
+ * 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.keyguard
+
+import com.android.systemui.Flags
+import com.android.systemui.flags.FlagToken
+import com.android.systemui.flags.RefactorFlagUtils
+
+/** Helper for reading or using the keyguard wm state refactor flag state. */
+@Suppress("NOTHING_TO_INLINE")
+object KeyguardWmStateRefactor {
+    /** The aconfig flag name */
+    const val FLAG_NAME = Flags.FLAG_KEYGUARD_WM_STATE_REFACTOR
+
+    /** A token used for dependency declaration */
+    val token: FlagToken
+        get() = FlagToken(FLAG_NAME, isEnabled)
+
+    /** Is the refactor enabled */
+    @JvmStatic
+    inline val isEnabled
+        get() = Flags.keyguardWmStateRefactor()
+
+    /**
+     * Called to ensure code is only run when the flag is enabled. This protects users from the
+     * unintended behaviors caused by accidentally running new logic, while also crashing on an eng
+     * build to ensure that the refactor author catches issues in testing.
+     */
+    @JvmStatic
+    inline fun isUnexpectedlyInLegacyMode() =
+        RefactorFlagUtils.isUnexpectedlyInLegacyMode(isEnabled, FLAG_NAME)
+
+    /**
+     * Called to ensure code is only run when the flag is disabled. This will throw an exception if
+     * the flag is enabled to ensure that the refactor author catches issues in testing.
+     */
+    @JvmStatic
+    inline fun assertInLegacyMode() = RefactorFlagUtils.assertInLegacyMode(isEnabled, FLAG_NAME)
+}
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromLockscreenTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromLockscreenTransitionInteractor.kt
index 8b2b45f..a4467e4a 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromLockscreenTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromLockscreenTransitionInteractor.kt
@@ -23,7 +23,7 @@
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.flags.FeatureFlags
-import com.android.systemui.flags.Flags
+import com.android.systemui.keyguard.KeyguardWmStateRefactor
 import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepository
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.shared.model.KeyguardSurfaceBehindModel
@@ -307,7 +307,7 @@
     }
 
     private fun listenForLockscreenToGone() {
-        if (flags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+        if (KeyguardWmStateRefactor.isEnabled) {
             return
         }
 
@@ -324,7 +324,7 @@
     }
 
     private fun listenForLockscreenToGoneDragging() {
-        if (flags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+        if (KeyguardWmStateRefactor.isEnabled) {
             return
         }
 
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromPrimaryBouncerTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromPrimaryBouncerTransitionInteractor.kt
index c62055f..561cb0a 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromPrimaryBouncerTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromPrimaryBouncerTransitionInteractor.kt
@@ -22,7 +22,7 @@
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.flags.FeatureFlags
-import com.android.systemui.flags.Flags
+import com.android.systemui.keyguard.KeyguardWmStateRefactor
 import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepository
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.shared.model.KeyguardSurfaceBehindModel
@@ -211,7 +211,7 @@
     }
 
     private fun listenForPrimaryBouncerToGone() {
-        if (flags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+        if (KeyguardWmStateRefactor.isEnabled) {
             // This is handled in KeyguardSecurityContainerController and
             // StatusBarKeyguardViewManager, which calls the transition interactor to kick off a
             // transition vs. listening to legacy state flags.
diff --git a/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java b/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java
index fd53423..fa4bfea 100644
--- a/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java
+++ b/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java
@@ -85,8 +85,8 @@
 import com.android.systemui.dagger.qualifiers.Main;
 import com.android.systemui.dump.DumpManager;
 import com.android.systemui.flags.FeatureFlags;
-import com.android.systemui.flags.Flags;
 import com.android.systemui.keyguard.KeyguardUnlockAnimationController;
+import com.android.systemui.keyguard.KeyguardWmStateRefactor;
 import com.android.systemui.keyguard.WakefulnessLifecycle;
 import com.android.systemui.keyguard.ui.view.InWindowLauncherUnlockAnimationManager;
 import com.android.systemui.model.SysUiState;
@@ -618,7 +618,7 @@
         mDisplayTracker = displayTracker;
         mUnfoldTransitionProgressForwarder = unfoldTransitionProgressForwarder;
 
-        if (!featureFlags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+        if (!KeyguardWmStateRefactor.isEnabled()) {
             mSysuiUnlockAnimationController = sysuiUnlockAnimationController;
         } else {
             mSysuiUnlockAnimationController = inWindowLauncherUnlockAnimationManager;
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManager.java
index 88347ab..4c83ca2 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManager.java
@@ -69,6 +69,7 @@
 import com.android.systemui.dreams.DreamOverlayStateController;
 import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.flags.Flags;
+import com.android.systemui.keyguard.KeyguardWmStateRefactor;
 import com.android.systemui.keyguard.domain.interactor.KeyguardDismissActionInteractor;
 import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor;
 import com.android.systemui.keyguard.domain.interactor.WindowManagerLockscreenVisibilityInteractor;
@@ -474,7 +475,7 @@
             mIsDocked = mDockManager.isDocked();
         }
 
-        if (mFlags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+        if (KeyguardWmStateRefactor.isEnabled()) {
             // Show the keyguard views whenever we've told WM that the lockscreen is visible.
             mShadeViewController.postToView(() ->
                     collectFlow(
@@ -1428,7 +1429,7 @@
             executeAfterKeyguardGoneAction();
         }
 
-        if (mFlags.isEnabled(Flags.KEYGUARD_WM_STATE_REFACTOR)) {
+        if (KeyguardWmStateRefactor.isEnabled()) {
             mKeyguardTransitionInteractor.startDismissKeyguardTransition();
         }
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java b/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java
index b57cf53..2593f47 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java
@@ -25,6 +25,7 @@
 import static com.android.internal.widget.LockPatternUtils.StrongAuthTracker.STRONG_AUTH_REQUIRED_AFTER_NON_STRONG_BIOMETRICS_TIMEOUT;
 import static com.android.internal.widget.LockPatternUtils.StrongAuthTracker.STRONG_AUTH_REQUIRED_AFTER_USER_LOCKDOWN;
 import static com.android.systemui.Flags.FLAG_REFACTOR_GET_CURRENT_USER;
+import static com.android.systemui.Flags.FLAG_KEYGUARD_WM_STATE_REFACTOR;
 import static com.android.systemui.keyguard.KeyguardViewMediator.DELAYED_KEYGUARD_ACTION;
 import static com.android.systemui.keyguard.KeyguardViewMediator.KEYGUARD_LOCK_AFTER_DELAY_DEFAULT;
 import static com.android.systemui.keyguard.KeyguardViewMediator.REBOOT_MAINLINE_UPDATE;
@@ -267,8 +268,8 @@
                 mUserTracker,
                 mSceneContainerFlags);
         mFeatureFlags = new FakeFeatureFlags();
-        mFeatureFlags.set(Flags.KEYGUARD_WM_STATE_REFACTOR, false);
         mSetFlagsRule.enableFlags(FLAG_REFACTOR_GET_CURRENT_USER);
+        mSetFlagsRule.disableFlags(FLAG_KEYGUARD_WM_STATE_REFACTOR);
 
         DejankUtils.setImmediate(true);
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionScenariosTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionScenariosTest.kt
index 8b6611f..655dbbd 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionScenariosTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionScenariosTest.kt
@@ -21,6 +21,7 @@
 import com.android.keyguard.KeyguardSecurityModel
 import com.android.keyguard.KeyguardSecurityModel.SecurityMode.PIN
 import com.android.systemui.Flags.FLAG_COMMUNAL_HUB
+import com.android.systemui.Flags.FLAG_KEYGUARD_WM_STATE_REFACTOR
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.bouncer.data.repository.FakeKeyguardBouncerRepository
 import com.android.systemui.communal.domain.interactor.CommunalInteractor
@@ -28,7 +29,6 @@
 import com.android.systemui.communal.shared.model.CommunalSceneKey
 import com.android.systemui.communal.shared.model.ObservableCommunalTransitionState
 import com.android.systemui.flags.FakeFeatureFlags
-import com.android.systemui.flags.Flags
 import com.android.systemui.keyguard.data.repository.FakeCommandQueue
 import com.android.systemui.keyguard.data.repository.FakeKeyguardRepository
 import com.android.systemui.keyguard.data.repository.FakeKeyguardSurfaceBehindRepository
@@ -133,8 +133,8 @@
 
         whenever(keyguardSecurityModel.getSecurityMode(anyInt())).thenReturn(PIN)
 
-        featureFlags = FakeFeatureFlags().apply { set(Flags.KEYGUARD_WM_STATE_REFACTOR, false) }
         mSetFlagsRule.enableFlags(FLAG_COMMUNAL_HUB)
+        featureFlags = FakeFeatureFlags()
 
         keyguardInteractor = createKeyguardInteractor()
 
@@ -286,6 +286,10 @@
                     transitionInteractor = transitionInteractor,
                 )
                 .apply { start() }
+
+        mSetFlagsRule.disableFlags(
+            FLAG_KEYGUARD_WM_STATE_REFACTOR,
+        )
     }
 
     @Test
diff --git a/packages/SystemUI/tests/src/com/android/systemui/recents/OverviewProxyServiceTest.kt b/packages/SystemUI/tests/src/com/android/systemui/recents/OverviewProxyServiceTest.kt
index 70a48f5..e9f2132 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/recents/OverviewProxyServiceTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/recents/OverviewProxyServiceTest.kt
@@ -131,7 +131,10 @@
         whenever(packageManager.resolveServiceAsUser(any(), anyInt(), anyInt()))
             .thenReturn(mock(ResolveInfo::class.java))
 
-        featureFlags.set(Flags.KEYGUARD_WM_STATE_REFACTOR, false)
+        mSetFlagsRule.disableFlags(
+            com.android.systemui.Flags.FLAG_KEYGUARD_WM_STATE_REFACTOR,
+        )
+
         subject =
             OverviewProxyService(
                 context,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManagerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManagerTest.java
index 8dde935..cb45315 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManagerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManagerTest.java
@@ -182,8 +182,10 @@
         when(mBouncerViewDelegate.getBackCallback()).thenReturn(mBouncerViewDelegateBackCallback);
         mFeatureFlags = new FakeFeatureFlags();
         mFeatureFlags.set(Flags.REFACTOR_KEYGUARD_DISMISS_INTENT, false);
-        mFeatureFlags.set(Flags.KEYGUARD_WM_STATE_REFACTOR, false);
-        mSetFlagsRule.disableFlags(com.android.systemui.Flags.FLAG_DEVICE_ENTRY_UDFPS_REFACTOR);
+        mSetFlagsRule.disableFlags(
+                com.android.systemui.Flags.FLAG_DEVICE_ENTRY_UDFPS_REFACTOR,
+                com.android.systemui.Flags.FLAG_KEYGUARD_WM_STATE_REFACTOR
+        );
 
         when(mNotificationShadeWindowController.getWindowRootView())
                 .thenReturn(mNotificationShadeWindowView);