Merge "Per-power group sleep and dim timeout" into main

10 files changed, 408 insertions, 28 deletions

diff --git a/core/api/system-current.txt b/core/api/system-current.txt
index 860089fb326f..3fde7497a090 100644
--- a/core/api/system-current.txt
+++ b/core/api/system-current.txt
@@ -3558,10 +3558,12 @@ package android.companion.virtual {
     method @Deprecated public int getDefaultActivityPolicy();
     method @Deprecated public int getDefaultNavigationPolicy();
     method public int getDevicePolicy(int);
+    method @FlaggedApi("android.companion.virtualdevice.flags.device_aware_display_power") @NonNull public java.time.Duration getDimDuration();
     method @FlaggedApi("android.companion.virtual.flags.vdm_custom_home") @Nullable public android.content.ComponentName getHomeComponent();
     method @FlaggedApi("android.companion.virtual.flags.vdm_custom_ime") @Nullable public android.content.ComponentName getInputMethodComponent();
     method public int getLockState();
     method @Nullable public String getName();
+    method @FlaggedApi("android.companion.virtualdevice.flags.device_aware_display_power") @NonNull public java.time.Duration getScreenOffTimeout();
     method @NonNull public java.util.Set<android.os.UserHandle> getUsersWithMatchingAccounts();
     method @NonNull public java.util.List<android.companion.virtual.sensor.VirtualSensorConfig> getVirtualSensorConfigs();
     method public void writeToParcel(@NonNull android.os.Parcel, int);
@@ -3595,10 +3597,12 @@ package android.companion.virtual {
     method @Deprecated @NonNull public android.companion.virtual.VirtualDeviceParams.Builder setBlockedActivities(@NonNull java.util.Set<android.content.ComponentName>);
     method @Deprecated @NonNull public android.companion.virtual.VirtualDeviceParams.Builder setBlockedCrossTaskNavigations(@NonNull java.util.Set<android.content.ComponentName>);
     method @NonNull public android.companion.virtual.VirtualDeviceParams.Builder setDevicePolicy(int, int);
+    method @FlaggedApi("android.companion.virtualdevice.flags.device_aware_display_power") @NonNull public android.companion.virtual.VirtualDeviceParams.Builder setDimDuration(@NonNull java.time.Duration);
     method @FlaggedApi("android.companion.virtual.flags.vdm_custom_home") @NonNull public android.companion.virtual.VirtualDeviceParams.Builder setHomeComponent(@Nullable android.content.ComponentName);
     method @FlaggedApi("android.companion.virtual.flags.vdm_custom_ime") @NonNull public android.companion.virtual.VirtualDeviceParams.Builder setInputMethodComponent(@Nullable android.content.ComponentName);
     method @NonNull @RequiresPermission(value=android.Manifest.permission.ADD_ALWAYS_UNLOCKED_DISPLAY, conditional=true) public android.companion.virtual.VirtualDeviceParams.Builder setLockState(int);
     method @NonNull public android.companion.virtual.VirtualDeviceParams.Builder setName(@NonNull String);
+    method @FlaggedApi("android.companion.virtualdevice.flags.device_aware_display_power") @NonNull public android.companion.virtual.VirtualDeviceParams.Builder setScreenOffTimeout(@NonNull java.time.Duration);
     method @NonNull public android.companion.virtual.VirtualDeviceParams.Builder setUsersWithMatchingAccounts(@NonNull java.util.Set<android.os.UserHandle>);
     method @NonNull public android.companion.virtual.VirtualDeviceParams.Builder setVirtualSensorCallback(@NonNull java.util.concurrent.Executor, @NonNull android.companion.virtual.sensor.VirtualSensorCallback);
     method @NonNull public android.companion.virtual.VirtualDeviceParams.Builder setVirtualSensorDirectChannelCallback(@NonNull java.util.concurrent.Executor, @NonNull android.companion.virtual.sensor.VirtualSensorDirectChannelCallback);
diff --git a/core/java/android/companion/virtual/VirtualDeviceParams.java b/core/java/android/companion/virtual/VirtualDeviceParams.java
index 6dad015cd63b..2be27dabcf90 100644
--- a/core/java/android/companion/virtual/VirtualDeviceParams.java
+++ b/core/java/android/companion/virtual/VirtualDeviceParams.java
@@ -333,6 +333,8 @@ public final class VirtualDeviceParams implements Parcelable {
     @Nullable private final IVirtualSensorCallback mVirtualSensorCallback;
     private final int mAudioPlaybackSessionId;
     private final int mAudioRecordingSessionId;
+    private final long mDimDuration;
+    private final long mScreenOffTimeout;
 
     private VirtualDeviceParams(
             @LockState int lockState,
@@ -348,7 +350,9 @@ public final class VirtualDeviceParams implements Parcelable {
             @NonNull List<VirtualSensorConfig> virtualSensorConfigs,
             @Nullable IVirtualSensorCallback virtualSensorCallback,
             int audioPlaybackSessionId,
-            int audioRecordingSessionId) {
+            int audioRecordingSessionId,
+            long dimDuration,
+            long screenOffTimeout) {
         mLockState = lockState;
         mUsersWithMatchingAccounts =
                 new ArraySet<>(Objects.requireNonNull(usersWithMatchingAccounts));
@@ -366,6 +370,8 @@ public final class VirtualDeviceParams implements Parcelable {
         mVirtualSensorCallback = virtualSensorCallback;
         mAudioPlaybackSessionId = audioPlaybackSessionId;
         mAudioRecordingSessionId = audioRecordingSessionId;
+        mDimDuration = dimDuration;
+        mScreenOffTimeout = screenOffTimeout;
     }
 
     @SuppressWarnings("unchecked")
@@ -386,6 +392,8 @@ public final class VirtualDeviceParams implements Parcelable {
         mAudioRecordingSessionId = parcel.readInt();
         mHomeComponent = parcel.readTypedObject(ComponentName.CREATOR);
         mInputMethodComponent = parcel.readTypedObject(ComponentName.CREATOR);
+        mDimDuration = parcel.readLong();
+        mScreenOffTimeout = parcel.readLong();
     }
 
     /**
@@ -397,6 +405,26 @@ public final class VirtualDeviceParams implements Parcelable {
     }
 
     /**
+     * Returns the dim duration for the displays of this device.
+     *
+     * @see Builder#setDimDuration(Duration)
+     */
+    @FlaggedApi(android.companion.virtualdevice.flags.Flags.FLAG_DEVICE_AWARE_DISPLAY_POWER)
+    public @NonNull Duration getDimDuration() {
+        return Duration.ofMillis(mDimDuration);
+    }
+
+    /**
+     * Returns the screen off timeout of the displays of this device.
+     *
+     * @see Builder#setDimDuration(Duration)
+     */
+    @FlaggedApi(android.companion.virtualdevice.flags.Flags.FLAG_DEVICE_AWARE_DISPLAY_POWER)
+    public @NonNull Duration getScreenOffTimeout() {
+        return Duration.ofMillis(mScreenOffTimeout);
+    }
+
+    /**
      * Returns the custom component used as home on all displays owned by this virtual device that
      * support home activities.
      *
@@ -619,6 +647,8 @@ public final class VirtualDeviceParams implements Parcelable {
         dest.writeInt(mAudioRecordingSessionId);
         dest.writeTypedObject(mHomeComponent, flags);
         dest.writeTypedObject(mInputMethodComponent, flags);
+        dest.writeLong(mDimDuration);
+        dest.writeLong(mScreenOffTimeout);
     }
 
     @Override
@@ -653,7 +683,9 @@ public final class VirtualDeviceParams implements Parcelable {
                 && Objects.equals(mHomeComponent, that.mHomeComponent)
                 && Objects.equals(mInputMethodComponent, that.mInputMethodComponent)
                 && mAudioPlaybackSessionId == that.mAudioPlaybackSessionId
-                && mAudioRecordingSessionId == that.mAudioRecordingSessionId;
+                && mAudioRecordingSessionId == that.mAudioRecordingSessionId
+                && mDimDuration == that.mDimDuration
+                && mScreenOffTimeout == that.mScreenOffTimeout;
     }
 
     @Override
@@ -662,7 +694,7 @@ public final class VirtualDeviceParams implements Parcelable {
                 mLockState, mUsersWithMatchingAccounts, mCrossTaskNavigationExemptions,
                 mDefaultNavigationPolicy, mActivityPolicyExemptions, mDefaultActivityPolicy, mName,
                 mDevicePolicies, mHomeComponent, mInputMethodComponent, mAudioPlaybackSessionId,
-                mAudioRecordingSessionId);
+                mAudioRecordingSessionId, mDimDuration, mScreenOffTimeout);
         for (int i = 0; i < mDevicePolicies.size(); i++) {
             hashCode = 31 * hashCode + mDevicePolicies.keyAt(i);
             hashCode = 31 * hashCode + mDevicePolicies.valueAt(i);
@@ -686,6 +718,8 @@ public final class VirtualDeviceParams implements Parcelable {
                 + " mInputMethodComponent=" + mInputMethodComponent
                 + " mAudioPlaybackSessionId=" + mAudioPlaybackSessionId
                 + " mAudioRecordingSessionId=" + mAudioRecordingSessionId
+                + " mDimDuration=" + mDimDuration
+                + " mScreenOffTimeout=" + mScreenOffTimeout
                 + ")";
     }
 
@@ -707,6 +741,8 @@ public final class VirtualDeviceParams implements Parcelable {
         pw.println(prefix + "mInputMethodComponent=" + mInputMethodComponent);
         pw.println(prefix + "mAudioPlaybackSessionId=" + mAudioPlaybackSessionId);
         pw.println(prefix + "mAudioRecordingSessionId=" + mAudioRecordingSessionId);
+        pw.println(prefix + "mDimDuration=" + mDimDuration);
+        pw.println(prefix + "mScreenOffTimeout=" + mScreenOffTimeout);
     }
 
     @NonNull
@@ -726,6 +762,8 @@ public final class VirtualDeviceParams implements Parcelable {
      */
     public static final class Builder {
 
+        private static final Duration INFINITE_TIMEOUT = Duration.ofDays(365 * 1000);
+
         private @LockState int mLockState = LOCK_STATE_DEFAULT;
         @NonNull private Set<UserHandle> mUsersWithMatchingAccounts = Collections.emptySet();
         @NonNull private Set<ComponentName> mCrossTaskNavigationExemptions = Collections.emptySet();
@@ -748,6 +786,8 @@ public final class VirtualDeviceParams implements Parcelable {
         @Nullable private VirtualSensorDirectChannelCallback mVirtualSensorDirectChannelCallback;
         @Nullable private ComponentName mHomeComponent;
         @Nullable private ComponentName mInputMethodComponent;
+        private Duration mDimDuration = Duration.ZERO;
+        private Duration mScreenOffTimeout = Duration.ZERO;
 
         private static class VirtualSensorCallbackDelegate extends IVirtualSensorCallback.Stub {
             @NonNull
@@ -825,6 +865,57 @@ public final class VirtualDeviceParams implements Parcelable {
         }
 
         /**
+         * Sets the dim duration for all trusted non-mirror displays of the device.
+         *
+         * <p>The system will reduce the display brightness for the specified duration if there
+         * has been no interaction just before the displays turn off.</p>
+         *
+         * <p>If set, the screen off timeout must also be set to a value larger than the dim
+         * duration. If left unset or set to zero, then the display brightness will not be reduced.
+         * </p>
+         *
+         * @throws IllegalArgumentException if the dim duration is negative or if the dim duration
+         *   is longer than the screen off timeout.
+         * @see android.hardware.display.DisplayManager#VIRTUAL_DISPLAY_FLAG_TRUSTED
+         * @see android.hardware.display.DisplayManager#VIRTUAL_DISPLAY_FLAG_AUTO_MIRROR
+         * @see #setScreenOffTimeout
+         */
+        @FlaggedApi(android.companion.virtualdevice.flags.Flags.FLAG_DEVICE_AWARE_DISPLAY_POWER)
+        @NonNull
+        public Builder setDimDuration(@NonNull Duration dimDuration) {
+            if (Objects.requireNonNull(dimDuration).compareTo(Duration.ZERO) < 0) {
+                throw new IllegalArgumentException("The dim duration cannot be negative");
+            }
+            mDimDuration = dimDuration;
+            return this;
+        }
+
+        /**
+         * Sets the timeout, after which all trusted non-mirror displays of the device will turn
+         * off, if there has been no interaction with the device.
+         *
+         * <p>If dim duration is set, the screen off timeout must be set to a value larger than the
+         * dim duration. If left unset or set to zero, then the displays will never be turned off
+         * due to inactivity.</p>
+         *
+         * @throws IllegalArgumentException if the screen off timeout is negative or if the dim
+         *   duration is longer than the screen off timeout.
+         * @see android.hardware.display.DisplayManager#VIRTUAL_DISPLAY_FLAG_TRUSTED
+         * @see android.hardware.display.DisplayManager#VIRTUAL_DISPLAY_FLAG_AUTO_MIRROR
+         * @see #setDimDuration
+         * @see VirtualDeviceManager.VirtualDevice#goToSleep()
+         */
+        @FlaggedApi(android.companion.virtualdevice.flags.Flags.FLAG_DEVICE_AWARE_DISPLAY_POWER)
+        @NonNull
+        public Builder setScreenOffTimeout(@NonNull Duration screenOffTimeout) {
+            if (Objects.requireNonNull(screenOffTimeout).compareTo(Duration.ZERO) < 0) {
+                throw new IllegalArgumentException("The screen off timeout cannot be negative");
+            }
+            mScreenOffTimeout = screenOffTimeout;
+            return this;
+        }
+
+        /**
          * Specifies a component to be used as home on all displays owned by this virtual device
          * that support home activities.
          * *
@@ -1220,6 +1311,14 @@ public final class VirtualDeviceParams implements Parcelable {
                 }
             }
 
+            if (mDimDuration.compareTo(mScreenOffTimeout) > 0) {
+                throw new IllegalArgumentException(
+                        "The dim duration cannot be greater than the screen off timeout.");
+            }
+            if (mScreenOffTimeout.compareTo(Duration.ZERO) == 0) {
+                mScreenOffTimeout = INFINITE_TIMEOUT;
+            }
+
             if (!Flags.crossDeviceClipboard()) {
                 mDevicePolicies.delete(POLICY_TYPE_CLIPBOARD);
             }
@@ -1269,7 +1368,9 @@ public final class VirtualDeviceParams implements Parcelable {
                     mVirtualSensorConfigs,
                     virtualSensorCallbackDelegate,
                     mAudioPlaybackSessionId,
-                    mAudioRecordingSessionId);
+                    mAudioRecordingSessionId,
+                    mDimDuration.toMillis(),
+                    mScreenOffTimeout.toMillis());
         }
     }
 }
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 281a2ce0556b..8b5b93e96494 100644
--- a/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java
+++ b/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java
@@ -1465,28 +1465,28 @@ final class VirtualDeviceImpl extends IVirtualDevice.Stub
     public int createVirtualDisplay(@NonNull VirtualDisplayConfig virtualDisplayConfig,
             @NonNull IVirtualDisplayCallback callback) {
         checkCallerIsDeviceOwner();
+
+        int displayId;
+        boolean showPointer;
+        boolean isTrustedDisplay;
         GenericWindowPolicyController gwpc;
         synchronized (mVirtualDeviceLock) {
             gwpc = createWindowPolicyControllerLocked(virtualDisplayConfig.getDisplayCategories());
-        }
-        int displayId;
-        displayId = mDisplayManagerInternal.createVirtualDisplay(virtualDisplayConfig, callback,
-                this, gwpc, mOwnerPackageName);
-        boolean isMirrorDisplay =
-                mDisplayManagerInternal.getDisplayIdToMirror(displayId) != Display.INVALID_DISPLAY;
-        gwpc.setDisplayId(displayId, isMirrorDisplay);
-        boolean isTrustedDisplay =
-                (mDisplayManagerInternal.getDisplayInfo(displayId).flags & Display.FLAG_TRUSTED)
-                        == Display.FLAG_TRUSTED;
-        if (!isTrustedDisplay) {
-            if (getDevicePolicy(POLICY_TYPE_CLIPBOARD) != DEVICE_POLICY_DEFAULT) {
-                throw new SecurityException("All displays must be trusted for devices with custom"
-                        + "clipboard policy.");
+            displayId = mDisplayManagerInternal.createVirtualDisplay(virtualDisplayConfig,
+                    callback, this, gwpc, mOwnerPackageName);
+            boolean isMirrorDisplay =
+                    mDisplayManagerInternal.getDisplayIdToMirror(displayId)
+                            != Display.INVALID_DISPLAY;
+            gwpc.setDisplayId(displayId, isMirrorDisplay);
+            isTrustedDisplay =
+                    (mDisplayManagerInternal.getDisplayInfo(displayId).flags & Display.FLAG_TRUSTED)
+                            == Display.FLAG_TRUSTED;
+            if (!isTrustedDisplay
+                    && getDevicePolicy(POLICY_TYPE_CLIPBOARD) != DEVICE_POLICY_DEFAULT) {
+                throw new SecurityException("All displays must be trusted for devices with "
+                        + "custom clipboard policy.");
             }
-        }
 
-        boolean showPointer;
-        synchronized (mVirtualDeviceLock) {
             if (mVirtualDisplays.contains(displayId)) {
                 gwpc.unregisterRunningAppsChangedListener(this);
                 throw new IllegalStateException(
@@ -1523,6 +1523,9 @@ final class VirtualDeviceImpl extends IVirtualDevice.Stub
     }
 
     private PowerManager.WakeLock createAndAcquireWakeLockForDisplay(int displayId) {
+        if (android.companion.virtualdevice.flags.Flags.deviceAwareDisplayPower()) {
+            return null;
+        }
         final long token = Binder.clearCallingIdentity();
         try {
             PowerManager powerManager = mContext.getSystemService(PowerManager.class);
@@ -1681,6 +1684,14 @@ final class VirtualDeviceImpl extends IVirtualDevice.Stub
         return mOwnerUid;
     }
 
+    long getDimDurationMillis() {
+        return mParams.getDimDuration().toMillis();
+    }
+
+    long getScreenOffTimeoutMillis() {
+        return mParams.getScreenOffTimeout().toMillis();
+    }
+
     @Override  // Binder call
     public int[] getDisplayIds() {
         synchronized (mVirtualDeviceLock) {
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 604aaa74306f..6729231d68ab 100644
--- a/services/companion/java/com/android/server/companion/virtual/VirtualDeviceManagerService.java
+++ b/services/companion/java/com/android/server/companion/virtual/VirtualDeviceManagerService.java
@@ -591,7 +591,6 @@ public class VirtualDeviceManagerService extends SystemService {
             }
         }
 
-
         @Override // Binder call
         public int getDeviceIdForDisplayId(int displayId) {
             if (displayId == Display.INVALID_DISPLAY || displayId == Display.DEFAULT_DISPLAY) {
@@ -926,6 +925,22 @@ public class VirtualDeviceManagerService extends SystemService {
         }
 
         @Override
+        public long getDimDurationMillisForDeviceId(int deviceId) {
+            synchronized (mVirtualDeviceManagerLock) {
+                VirtualDeviceImpl virtualDevice = mVirtualDevices.get(deviceId);
+                return virtualDevice == null ? -1 : virtualDevice.getDimDurationMillis();
+            }
+        }
+
+        @Override
+        public long getScreenOffTimeoutMillisForDeviceId(int deviceId) {
+            synchronized (mVirtualDeviceManagerLock) {
+                VirtualDeviceImpl virtualDevice = mVirtualDevices.get(deviceId);
+                return virtualDevice == null ? -1 : virtualDevice.getScreenOffTimeoutMillis();
+            }
+        }
+
+        @Override
         public boolean isValidVirtualDeviceId(int deviceId) {
             return mImpl.isValidVirtualDeviceId(deviceId);
         }
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 6e38733f04c2..471b7b4ddfc8 100644
--- a/services/core/java/com/android/server/companion/virtual/VirtualDeviceManagerInternal.java
+++ b/services/core/java/com/android/server/companion/virtual/VirtualDeviceManagerInternal.java
@@ -167,6 +167,12 @@ public abstract class VirtualDeviceManagerInternal {
      */
     public abstract int getDeviceIdForDisplayId(int displayId);
 
+    /** Returns the dim duration for the displays of the device with the given ID. */
+    public abstract long getDimDurationMillisForDeviceId(int deviceId);
+
+    /** Returns the screen off timeout of the displays of the device with the given ID. */
+    public abstract long getScreenOffTimeoutMillisForDeviceId(int deviceId);
+
     /**
      * Gets the persistent ID for the VirtualDevice with the given device ID.
      *
diff --git a/services/core/java/com/android/server/power/PowerGroup.java b/services/core/java/com/android/server/power/PowerGroup.java
index a928814c7909..01a2045df426 100644
--- a/services/core/java/com/android/server/power/PowerGroup.java
+++ b/services/core/java/com/android/server/power/PowerGroup.java
@@ -42,6 +42,8 @@ import android.view.Display;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.util.LatencyTracker;
+import com.android.server.LocalServices;
+import com.android.server.companion.virtual.VirtualDeviceManagerInternal;
 import com.android.server.power.feature.PowerManagerFlags;
 
 /**
@@ -56,6 +58,11 @@ public class PowerGroup {
     private static final String TAG = PowerGroup.class.getSimpleName();
     private static final boolean DEBUG = false;
 
+    /**
+     * Indicates that the default dim/sleep timeouts should be used.
+     */
+    private static final long INVALID_TIMEOUT = -1;
+
     @VisibleForTesting
     final DisplayPowerRequest mDisplayPowerRequest = new DisplayPowerRequest();
     private final PowerGroupListener mWakefulnessListener;
@@ -91,6 +98,9 @@ public class PowerGroup {
     private @PowerManager.GoToSleepReason int mLastSleepReason =
             PowerManager.GO_TO_SLEEP_REASON_UNKNOWN;
 
+    private final long mDimDuration;
+    private final long mScreenOffTimeout;
+
     PowerGroup(int groupId, PowerGroupListener wakefulnessListener, Notifier notifier,
             DisplayManagerInternal displayManagerInternal, int wakefulness, boolean ready,
             boolean supportsSandman, long eventTime, PowerManagerFlags featureFlags) {
@@ -104,6 +114,30 @@ public class PowerGroup {
         mLastWakeTime = eventTime;
         mLastSleepTime = eventTime;
         mFeatureFlags = featureFlags;
+
+        long dimDuration = INVALID_TIMEOUT;
+        long screenOffTimeout = INVALID_TIMEOUT;
+        if (android.companion.virtualdevice.flags.Flags.deviceAwareDisplayPower()
+                && mGroupId != Display.DEFAULT_DISPLAY_GROUP) {
+            VirtualDeviceManagerInternal vdm =
+                    LocalServices.getService(VirtualDeviceManagerInternal.class);
+            if (vdm != null) {
+                int[] displayIds = mDisplayManagerInternal.getDisplayIdsForGroup(mGroupId);
+                if (displayIds != null && displayIds.length > 0) {
+                    int deviceId = vdm.getDeviceIdForDisplayId(displayIds[0]);
+                    if (vdm.isValidVirtualDeviceId(deviceId)) {
+                        dimDuration = vdm.getDimDurationMillisForDeviceId(deviceId);
+                        screenOffTimeout = vdm.getScreenOffTimeoutMillisForDeviceId(deviceId);
+                        if (dimDuration > 0 && dimDuration > screenOffTimeout) {
+                            // If the dim duration is set, cap it to the screen off timeout.
+                            dimDuration = screenOffTimeout;
+                        }
+                    }
+                }
+            }
+        }
+        mDimDuration = dimDuration;
+        mScreenOffTimeout = screenOffTimeout;
     }
 
     PowerGroup(int wakefulness, PowerGroupListener wakefulnessListener, Notifier notifier,
@@ -119,6 +153,16 @@ public class PowerGroup {
         mLastWakeTime = eventTime;
         mLastSleepTime = eventTime;
         mFeatureFlags = featureFlags;
+        mDimDuration = INVALID_TIMEOUT;
+        mScreenOffTimeout = INVALID_TIMEOUT;
+    }
+
+    long getScreenOffTimeoutOverrideLocked(long defaultScreenOffTimeout) {
+        return mScreenOffTimeout == INVALID_TIMEOUT ? defaultScreenOffTimeout : mScreenOffTimeout;
+    }
+
+    long getScreenDimDurationOverrideLocked(long defaultScreenDimDuration) {
+        return mDimDuration == INVALID_TIMEOUT ? defaultScreenDimDuration : mDimDuration;
     }
 
     long getLastWakeTimeLocked() {
diff --git a/services/core/java/com/android/server/power/PowerManagerService.java b/services/core/java/com/android/server/power/PowerManagerService.java
index 3a5afacd0977..0acfe92f578d 100644
--- a/services/core/java/com/android/server/power/PowerManagerService.java
+++ b/services/core/java/com/android/server/power/PowerManagerService.java
@@ -2971,8 +2971,8 @@ public final class PowerManagerService extends SystemService
         mHandler.removeMessages(MSG_USER_ACTIVITY_TIMEOUT);
 
         final long attentiveTimeout = getAttentiveTimeoutLocked();
-        final long sleepTimeout = getSleepTimeoutLocked(attentiveTimeout);
-        final long defaultScreenOffTimeout = getScreenOffTimeoutLocked(sleepTimeout,
+        final long defaultSleepTimeout = getSleepTimeoutLocked(attentiveTimeout);
+        final long defaultScreenOffTimeout = getScreenOffTimeoutLocked(defaultSleepTimeout,
                 attentiveTimeout);
         final long defaultScreenDimDuration = getScreenDimDurationLocked(defaultScreenOffTimeout);
 
@@ -2985,13 +2985,25 @@ public final class PowerManagerService extends SystemService
             final PowerGroup powerGroup = mPowerGroups.valueAt(idx);
             final int wakefulness = powerGroup.getWakefulnessLocked();
 
-            // The default display screen timeout could be overridden by policy.
+            // The timeouts could be overridden by the power group policy.
             long screenOffTimeout = defaultScreenOffTimeout;
             long screenDimDuration = defaultScreenDimDuration;
+            long sleepTimeout = defaultSleepTimeout;
+            // TODO(b/376211497): Consolidate the timeout logic for all power groups.
             if (powerGroup.getGroupId() == Display.DEFAULT_DISPLAY_GROUP) {
                 screenOffTimeout =
-                        getScreenOffTimeoutOverrideLocked(screenOffTimeout, screenDimDuration);
+                        getDefaultGroupScreenOffTimeoutOverrideLocked(screenOffTimeout,
+                                screenDimDuration);
                 screenDimDuration = getScreenDimDurationLocked(screenOffTimeout);
+            } else {
+                screenOffTimeout = powerGroup.getScreenOffTimeoutOverrideLocked(screenOffTimeout);
+                screenDimDuration =
+                        powerGroup.getScreenDimDurationOverrideLocked(screenDimDuration);
+                if (sleepTimeout > 0 && screenOffTimeout > 0) {
+                    // If both sleep and screen off timeouts are set, make sure that the sleep
+                    // timeout is not smaller than the screen off one.
+                    sleepTimeout = Math.max(sleepTimeout, screenOffTimeout);
+                }
             }
 
             if (wakefulness != WAKEFULNESS_ASLEEP) {
@@ -3273,7 +3285,8 @@ public final class PowerManagerService extends SystemService
 
     @VisibleForTesting
     @GuardedBy("mLock")
-    long getScreenOffTimeoutOverrideLocked(long screenOffTimeout, long screenDimDuration) {
+    long getDefaultGroupScreenOffTimeoutOverrideLocked(long screenOffTimeout,
+            long screenDimDuration) {
         long shortestScreenOffTimeout = screenOffTimeout;
         if (mScreenTimeoutOverridePolicy != null) {
             shortestScreenOffTimeout =
diff --git a/services/tests/powerservicetests/src/com/android/server/power/PowerGroupTest.java b/services/tests/powerservicetests/src/com/android/server/power/PowerGroupTest.java
index 648da6530eb0..4e29e74651b6 100644
--- a/services/tests/powerservicetests/src/com/android/server/power/PowerGroupTest.java
+++ b/services/tests/powerservicetests/src/com/android/server/power/PowerGroupTest.java
@@ -49,17 +49,23 @@ import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.content.Context;
 import android.hardware.display.DisplayManagerInternal;
 import android.os.PowerManager;
 import android.os.PowerSaveState;
+import android.platform.test.annotations.EnableFlags;
+import android.platform.test.flag.junit.SetFlagsRule;
 import android.view.Display;
 
 import androidx.test.platform.app.InstrumentationRegistry;
 
 import com.android.internal.util.LatencyTracker;
+import com.android.server.LocalServices;
+import com.android.server.companion.virtual.VirtualDeviceManagerInternal;
 import com.android.server.power.feature.PowerManagerFlags;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
@@ -73,6 +79,9 @@ import org.mockito.MockitoAnnotations;
 public class PowerGroupTest {
 
     private static final int GROUP_ID = 0;
+    private static final int NON_DEFAULT_GROUP_ID = 1;
+    private static final int NON_DEFAULT_DISPLAY_ID = 2;
+    private static final int VIRTUAL_DEVICE_ID = 3;
     private static final int UID = 11;
     private static final long TIMESTAMP_CREATE = 1;
     private static final long TIMESTAMP1 = 999;
@@ -87,10 +96,16 @@ public class PowerGroupTest {
     private static final LatencyTracker LATENCY_TRACKER = LatencyTracker.getInstance(
             InstrumentationRegistry.getInstrumentation().getContext());
 
+    private static final long DEFAULT_TIMEOUT = 1234L;
+
+    @Rule
+    public SetFlagsRule mSetFlagsRule = new SetFlagsRule();
+
     private PowerGroup mPowerGroup;
     @Mock private PowerGroup.PowerGroupListener mWakefulnessCallbackMock;
     @Mock private Notifier mNotifier;
     @Mock private DisplayManagerInternal mDisplayManagerInternal;
+    @Mock private VirtualDeviceManagerInternal mVirtualDeviceManagerInternal;
     @Mock private PowerManagerFlags mFeatureFlags;
 
 
@@ -740,4 +755,111 @@ public class PowerGroupTest {
         assertThat(displayPowerRequest.screenLowPowerBrightnessFactor).isWithin(PRECISION).of(
                 brightnessFactor);
     }
+
+    @Test
+    public void testTimeoutsOverride_defaultGroup_noOverride() {
+        assertThat(mPowerGroup.getScreenDimDurationOverrideLocked(DEFAULT_TIMEOUT))
+                .isEqualTo(DEFAULT_TIMEOUT);
+        assertThat(mPowerGroup.getScreenOffTimeoutOverrideLocked(DEFAULT_TIMEOUT))
+                .isEqualTo(DEFAULT_TIMEOUT);
+    }
+
+    @EnableFlags(android.companion.virtualdevice.flags.Flags.FLAG_DEVICE_AWARE_DISPLAY_POWER)
+    @Test
+    public void testTimeoutsOverride_noVdm_noOverride() {
+        LocalServices.removeServiceForTest(VirtualDeviceManagerInternal.class);
+        LocalServices.addService(VirtualDeviceManagerInternal.class, null);
+
+        mPowerGroup = new PowerGroup(NON_DEFAULT_GROUP_ID, mWakefulnessCallbackMock, mNotifier,
+                mDisplayManagerInternal, WAKEFULNESS_AWAKE, /* ready= */ true,
+                /* supportsSandman= */ true, TIMESTAMP_CREATE, mFeatureFlags);
+
+        assertThat(mPowerGroup.getScreenDimDurationOverrideLocked(DEFAULT_TIMEOUT))
+                .isEqualTo(DEFAULT_TIMEOUT);
+        assertThat(mPowerGroup.getScreenOffTimeoutOverrideLocked(DEFAULT_TIMEOUT))
+                .isEqualTo(DEFAULT_TIMEOUT);
+    }
+
+    @EnableFlags(android.companion.virtualdevice.flags.Flags.FLAG_DEVICE_AWARE_DISPLAY_POWER)
+    @Test
+    public void testTimeoutsOverride_notValidVirtualDeviceId_noOverride() {
+        LocalServices.removeServiceForTest(VirtualDeviceManagerInternal.class);
+        LocalServices.addService(VirtualDeviceManagerInternal.class, mVirtualDeviceManagerInternal);
+
+        when(mDisplayManagerInternal.getDisplayIdsForGroup(NON_DEFAULT_GROUP_ID))
+                .thenReturn(new int[] {NON_DEFAULT_DISPLAY_ID});
+        when(mVirtualDeviceManagerInternal.getDeviceIdForDisplayId(NON_DEFAULT_DISPLAY_ID))
+                .thenReturn(Context.DEVICE_ID_DEFAULT);
+        when(mVirtualDeviceManagerInternal.isValidVirtualDeviceId(Context.DEVICE_ID_DEFAULT))
+                .thenReturn(false);
+
+        mPowerGroup = new PowerGroup(NON_DEFAULT_GROUP_ID, mWakefulnessCallbackMock, mNotifier,
+                mDisplayManagerInternal, WAKEFULNESS_AWAKE, /* ready= */ true,
+                /* supportsSandman= */ true, TIMESTAMP_CREATE, mFeatureFlags);
+
+        assertThat(mPowerGroup.getScreenDimDurationOverrideLocked(DEFAULT_TIMEOUT))
+                .isEqualTo(DEFAULT_TIMEOUT);
+        assertThat(mPowerGroup.getScreenOffTimeoutOverrideLocked(DEFAULT_TIMEOUT))
+                .isEqualTo(DEFAULT_TIMEOUT);
+    }
+
+    @EnableFlags(android.companion.virtualdevice.flags.Flags.FLAG_DEVICE_AWARE_DISPLAY_POWER)
+    @Test
+    public void testTimeoutsOverride_validVirtualDeviceId_timeoutsAreOverridden() {
+        LocalServices.removeServiceForTest(VirtualDeviceManagerInternal.class);
+        LocalServices.addService(VirtualDeviceManagerInternal.class, mVirtualDeviceManagerInternal);
+
+        final long dimDurationOverride = DEFAULT_TIMEOUT * 3;
+        final long screenOffTimeoutOverride = DEFAULT_TIMEOUT * 5;
+
+        when(mDisplayManagerInternal.getDisplayIdsForGroup(NON_DEFAULT_GROUP_ID))
+                .thenReturn(new int[] {NON_DEFAULT_DISPLAY_ID});
+        when(mVirtualDeviceManagerInternal.getDeviceIdForDisplayId(NON_DEFAULT_DISPLAY_ID))
+                .thenReturn(VIRTUAL_DEVICE_ID);
+        when(mVirtualDeviceManagerInternal.isValidVirtualDeviceId(VIRTUAL_DEVICE_ID))
+                .thenReturn(true);
+        when(mVirtualDeviceManagerInternal.getDimDurationMillisForDeviceId(VIRTUAL_DEVICE_ID))
+                .thenReturn(dimDurationOverride);
+        when(mVirtualDeviceManagerInternal.getScreenOffTimeoutMillisForDeviceId(VIRTUAL_DEVICE_ID))
+                .thenReturn(screenOffTimeoutOverride);
+
+        mPowerGroup = new PowerGroup(NON_DEFAULT_GROUP_ID, mWakefulnessCallbackMock, mNotifier,
+                mDisplayManagerInternal, WAKEFULNESS_AWAKE, /* ready= */ true,
+                /* supportsSandman= */ true, TIMESTAMP_CREATE, mFeatureFlags);
+
+        assertThat(mPowerGroup.getScreenDimDurationOverrideLocked(DEFAULT_TIMEOUT))
+                .isEqualTo(dimDurationOverride);
+        assertThat(mPowerGroup.getScreenOffTimeoutOverrideLocked(DEFAULT_TIMEOUT))
+                .isEqualTo(screenOffTimeoutOverride);
+    }
+
+    @EnableFlags(android.companion.virtualdevice.flags.Flags.FLAG_DEVICE_AWARE_DISPLAY_POWER)
+    @Test
+    public void testTimeoutsOverrides_dimDurationIsCapped() {
+        LocalServices.removeServiceForTest(VirtualDeviceManagerInternal.class);
+        LocalServices.addService(VirtualDeviceManagerInternal.class, mVirtualDeviceManagerInternal);
+
+        final long dimDurationOverride = DEFAULT_TIMEOUT * 5;
+        final long screenOffTimeoutOverride = DEFAULT_TIMEOUT * 3;
+
+        when(mDisplayManagerInternal.getDisplayIdsForGroup(NON_DEFAULT_GROUP_ID))
+                .thenReturn(new int[] {NON_DEFAULT_DISPLAY_ID});
+        when(mVirtualDeviceManagerInternal.getDeviceIdForDisplayId(NON_DEFAULT_DISPLAY_ID))
+                .thenReturn(VIRTUAL_DEVICE_ID);
+        when(mVirtualDeviceManagerInternal.isValidVirtualDeviceId(VIRTUAL_DEVICE_ID))
+                .thenReturn(true);
+        when(mVirtualDeviceManagerInternal.getDimDurationMillisForDeviceId(VIRTUAL_DEVICE_ID))
+                .thenReturn(dimDurationOverride);
+        when(mVirtualDeviceManagerInternal.getScreenOffTimeoutMillisForDeviceId(VIRTUAL_DEVICE_ID))
+                .thenReturn(screenOffTimeoutOverride);
+
+        mPowerGroup = new PowerGroup(NON_DEFAULT_GROUP_ID, mWakefulnessCallbackMock, mNotifier,
+                mDisplayManagerInternal, WAKEFULNESS_AWAKE, /* ready= */ true,
+                /* supportsSandman= */ true, TIMESTAMP_CREATE, mFeatureFlags);
+
+        assertThat(mPowerGroup.getScreenDimDurationOverrideLocked(DEFAULT_TIMEOUT))
+                .isEqualTo(screenOffTimeoutOverride);
+        assertThat(mPowerGroup.getScreenOffTimeoutOverrideLocked(DEFAULT_TIMEOUT))
+                .isEqualTo(screenOffTimeoutOverride);
+    }
 }
diff --git a/services/tests/powerservicetests/src/com/android/server/power/PowerManagerServiceTest.java b/services/tests/powerservicetests/src/com/android/server/power/PowerManagerServiceTest.java
index b10200daeeb4..359cf6376239 100644
--- a/services/tests/powerservicetests/src/com/android/server/power/PowerManagerServiceTest.java
+++ b/services/tests/powerservicetests/src/com/android/server/power/PowerManagerServiceTest.java
@@ -85,6 +85,7 @@ import android.os.PowerManagerInternal;
 import android.os.PowerSaveState;
 import android.os.UserHandle;
 import android.os.test.TestLooper;
+import android.platform.test.annotations.EnableFlags;
 import android.platform.test.flag.junit.SetFlagsRule;
 import android.provider.DeviceConfig;
 import android.provider.Settings;
@@ -102,6 +103,7 @@ import com.android.internal.foldables.FoldGracePeriodProvider;
 import com.android.internal.util.test.FakeSettingsProvider;
 import com.android.server.LocalServices;
 import com.android.server.SystemService;
+import com.android.server.companion.virtual.VirtualDeviceManagerInternal;
 import com.android.server.display.feature.DeviceConfigParameterProvider;
 import com.android.server.lights.LightsManager;
 import com.android.server.policy.WindowManagerPolicy;
@@ -167,6 +169,7 @@ public class PowerManagerServiceTest {
     @Mock private BatteryManagerInternal mBatteryManagerInternalMock;
     @Mock private ActivityManagerInternal mActivityManagerInternalMock;
     @Mock private AttentionManagerInternal mAttentionManagerInternalMock;
+    @Mock private VirtualDeviceManagerInternal mVirtualDeviceManagerInternalMock;
     @Mock private DreamManagerInternal mDreamManagerInternalMock;
     @Mock private PowerManagerService.NativeWrapper mNativeWrapperMock;
     @Mock private FoldGracePeriodProvider mFoldGracePeriodProvider;
@@ -246,6 +249,7 @@ public class PowerManagerServiceTest {
         addLocalServiceMock(ActivityManagerInternal.class, mActivityManagerInternalMock);
         addLocalServiceMock(AttentionManagerInternal.class, mAttentionManagerInternalMock);
         addLocalServiceMock(DreamManagerInternal.class, mDreamManagerInternalMock);
+        addLocalServiceMock(VirtualDeviceManagerInternal.class, mVirtualDeviceManagerInternalMock);
 
         mContextSpy = spy(new ContextWrapper(ApplicationProvider.getApplicationContext()));
         mResourcesSpy = spy(mContextSpy.getResources());
@@ -1200,6 +1204,59 @@ public class PowerManagerServiceTest {
         assertThat(mService.getGlobalWakefulnessLocked()).isEqualTo(WAKEFULNESS_DREAMING);
     }
 
+    @EnableFlags(android.companion.virtualdevice.flags.Flags.FLAG_DEVICE_AWARE_DISPLAY_POWER)
+    @SuppressWarnings("GuardedBy")
+    @Test
+    public void testNonDefaultDisplayGroupWithCustomTimeout_afterTimeout_goesToDozing() {
+        final int nonDefaultDisplayGroupId = Display.DEFAULT_DISPLAY_GROUP + 1;
+        final int nonDefaultDisplayId = Display.DEFAULT_DISPLAY + 2;
+        final int virtualDeviceId = Context.DEVICE_ID_DEFAULT + 3;
+        final AtomicReference<DisplayManagerInternal.DisplayGroupListener> listener =
+                new AtomicReference<>();
+        doAnswer((Answer<Void>) invocation -> {
+            listener.set(invocation.getArgument(0));
+            return null;
+        }).when(mDisplayManagerInternalMock).registerDisplayGroupListener(any());
+        final DisplayInfo info = new DisplayInfo();
+        info.displayGroupId = nonDefaultDisplayGroupId;
+        when(mDisplayManagerInternalMock.getDisplayInfo(nonDefaultDisplayId)).thenReturn(info);
+        when(mDisplayManagerInternalMock.getDisplayIdsForGroup(nonDefaultDisplayGroupId))
+                .thenReturn(new int[] {nonDefaultDisplayId});
+        when(mVirtualDeviceManagerInternalMock.getDeviceIdForDisplayId(nonDefaultDisplayId))
+                .thenReturn(virtualDeviceId);
+        when(mVirtualDeviceManagerInternalMock.isValidVirtualDeviceId(virtualDeviceId))
+                .thenReturn(true);
+        when(mVirtualDeviceManagerInternalMock
+                .getScreenOffTimeoutMillisForDeviceId(virtualDeviceId))
+                .thenReturn(20000L);
+
+        setMinimumScreenOffTimeoutConfig(10000);
+        createService();
+        startSystem();
+        listener.get().onDisplayGroupAdded(nonDefaultDisplayGroupId);
+
+        assertThat(mService.getWakefulnessLocked(Display.DEFAULT_DISPLAY_GROUP))
+                .isEqualTo(WAKEFULNESS_AWAKE);
+        assertThat(mService.getWakefulnessLocked(nonDefaultDisplayGroupId))
+                .isEqualTo(WAKEFULNESS_AWAKE);
+
+        // The default timeout is 10s, the custom group timeout is 20s.
+
+        advanceTime(15000);
+
+        assertThat(mService.getWakefulnessLocked(Display.DEFAULT_DISPLAY_GROUP))
+                .isEqualTo(WAKEFULNESS_ASLEEP);
+        assertThat(mService.getWakefulnessLocked(nonDefaultDisplayGroupId))
+                .isEqualTo(WAKEFULNESS_AWAKE);
+
+        advanceTime(10000);
+
+        assertThat(mService.getWakefulnessLocked(Display.DEFAULT_DISPLAY_GROUP))
+                .isEqualTo(WAKEFULNESS_ASLEEP);
+        assertThat(mService.getWakefulnessLocked(nonDefaultDisplayGroupId))
+                .isEqualTo(WAKEFULNESS_DOZING);
+    }
+
     @SuppressWarnings("GuardedBy")
     @Test
     public void testAmbientSuppression_disablesDreamingAndWakesDevice() {
@@ -3348,7 +3405,8 @@ public class PowerManagerServiceTest {
                         null /* workSource */, null /* historyTag */, Display.INVALID_DISPLAY,
                         null /* callback */);
             }
-            assertThat(mService.getScreenOffTimeoutOverrideLocked(screenTimeout, screenDimTimeout))
+            assertThat(mService.getDefaultGroupScreenOffTimeoutOverrideLocked(screenTimeout,
+                           screenDimTimeout))
                     .isEqualTo(expect[2]);
             if (acquireWakeLock) {
                 mService.getBinderServiceInstance().releaseWakeLock(token, 0);
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 e3868089b21f..727d1b59646a 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
@@ -64,6 +64,7 @@ import android.companion.virtual.audio.IAudioRoutingCallback;
 import android.companion.virtual.sensor.VirtualSensor;
 import android.companion.virtual.sensor.VirtualSensorCallback;
 import android.companion.virtual.sensor.VirtualSensorConfig;
+import android.companion.virtualdevice.flags.Flags;
 import android.content.AttributionSource;
 import android.content.ComponentName;
 import android.content.Context;
@@ -103,6 +104,7 @@ import android.os.PowerManager;
 import android.os.Process;
 import android.os.RemoteException;
 import android.os.WorkSource;
+import android.platform.test.annotations.DisableFlags;
 import android.platform.test.annotations.Presubmit;
 import android.platform.test.flag.junit.SetFlagsRule;
 import android.testing.AndroidTestingRunner;
@@ -999,6 +1001,7 @@ public class VirtualDeviceManagerServiceTest {
                 nullable(String.class), anyInt(), eq(null));
     }
 
+    @DisableFlags(Flags.FLAG_DEVICE_AWARE_DISPLAY_POWER)
     @Test
     public void onVirtualDisplayCreatedLocked_wakeLockIsAcquired() throws RemoteException {
         verify(mIPowerManagerMock, never()).acquireWakeLock(any(Binder.class), anyInt(),
@@ -1010,6 +1013,7 @@ public class VirtualDeviceManagerServiceTest {
                 nullable(String.class), eq(DISPLAY_ID_1), eq(null));
     }
 
+    @DisableFlags(Flags.FLAG_DEVICE_AWARE_DISPLAY_POWER)
     @Test
     public void onVirtualDisplayCreatedLocked_duplicateCalls_onlyOneWakeLockIsAcquired()
             throws RemoteException {
@@ -1022,6 +1026,7 @@ public class VirtualDeviceManagerServiceTest {
                 nullable(String.class), eq(DISPLAY_ID_1), eq(null));
     }
 
+    @DisableFlags(Flags.FLAG_DEVICE_AWARE_DISPLAY_POWER)
     @Test
     public void onVirtualDisplayRemovedLocked_wakeLockIsReleased() throws RemoteException {
         addVirtualDisplay(mDeviceImpl, DISPLAY_ID_1, Display.FLAG_TRUSTED);
@@ -1037,6 +1042,7 @@ public class VirtualDeviceManagerServiceTest {
         verify(mIPowerManagerMock).releaseWakeLock(eq(wakeLock), anyInt());
     }
 
+    @DisableFlags(Flags.FLAG_DEVICE_AWARE_DISPLAY_POWER)
     @Test
     public void addVirtualDisplay_displayNotReleased_wakeLockIsReleased() throws RemoteException {
         addVirtualDisplay(mDeviceImpl, DISPLAY_ID_1, Display.FLAG_TRUSTED);