Awake Updates - TILLSON_11272024 (#36049)

* Update with bug fixes for tray icon and support for parent process

* Process information enum

* Update the docs

* Fix spelling

* Make sure that PID is used in PT config flow

* Logic for checks based on #34148

* Update with link to PR

* Fixes #34717

* Small cleanup

* Proper task segmentation in a function

* Cleanup the code

* Fix synchronization context issue

* Update planning doc

* Test disabling caching to see if that manages to pass CI

* Cleanup to make sure that we're logging things properly.

* Update ci.yml

* Disable cache to pass CI

* Retry logic

* Cleanup

* Code cleanup

* Fixes #35848

* Update notes and codename

* After third attempt, log error instead of throwing exception

* More cleanup to avoid double execution

* Add expected word

* Safeguards for bad values for timed keep-awake

* More updates to make sure I am using uint

* Update error message

* Update packages

* Fix notice and revert CsWinRT upgrade

* Codename update

* Update expect.txt

* Update the struct

* Ensuring we're properly awaiting tray initialization

* Update to make sure tray reflects the bound process

* Cleanup, proper JSON serialization for logs.

* Not needed.

* Add command validation logic

* Moving the initialization logic earlier

* Make sure we show the display state in the tooltip

* Update tray string

* Update src/modules/awake/Awake/Core/Manager.cs

Co-authored-by: Jaime Bernardo <jaime@janeasystems.com>

* Update src/modules/awake/Awake/Core/Manager.cs

Co-authored-by: Jaime Bernardo <jaime@janeasystems.com>

* Update src/modules/awake/Awake/Core/Manager.cs

Co-authored-by: Jaime Bernardo <jaime@janeasystems.com>

* Update src/modules/awake/Awake/Core/Manager.cs

Co-authored-by: Jaime Bernardo <jaime@janeasystems.com>

* Update logic for icon resets

* Update doc

* Simplify function for setting mode shell icon

* Issues should be properly linked

* Minor cleanup

* Update timed behavior

---------

Co-authored-by: Jaime Bernardo <jaime@janeasystems.com>
Co-authored-by: Clint Rutkas <clint@rutkas.com>

.github/actions/spell-check/expect.txt

@@ -1542,6 +1542,7 @@ THISCOMPONENT
 THotkey
 thumbcache
 TILEDWINDOW
+TILLSON
 timedate
 timediff
 timeunion

Directory.Packages.props

@@ -26,7 +26,7 @@
     <PackageVersion Include="Markdig.Signed" Version="0.34.0" />
     <!-- Including MessagePack to force version, since it's used by StreamJsonRpc but contains vulnerabilities. After StreamJsonRpc updates the version of MessagePack, we can upgrade StreamJsonRpc instead. -->
     <PackageVersion Include="MessagePack" Version="2.5.187" />
-    <PackageVersion Include="Microsoft.CodeAnalysis.NetAnalyzers" Version="9.0.0-preview.24508.2" />
+    <PackageVersion Include="Microsoft.CodeAnalysis.NetAnalyzers" Version="9.0.0" />
     <PackageVersion Include="Microsoft.Data.Sqlite" Version="9.0.0" />
     <PackageVersion Include="Microsoft.Diagnostics.Tracing.TraceEvent" Version="3.1.16" />
     <PackageVersion Include="Microsoft.Extensions.DependencyInjection" Version="9.0.0" />

NOTICE.md

@@ -1318,7 +1318,7 @@ EXHIBIT A -Mozilla Public License.
 - Mages 2.0.2
 - Markdig.Signed 0.34.0
 - MessagePack 2.5.187
-- Microsoft.CodeAnalysis.NetAnalyzers 9.0.0-preview.24508.2
+- Microsoft.CodeAnalysis.NetAnalyzers 9.0.0
 - Microsoft.Data.Sqlite 9.0.0
 - Microsoft.Diagnostics.Tracing.TraceEvent 3.1.16
 - Microsoft.Extensions.DependencyInjection 9.0.0

doc/planning/awake.md

@@ -10,8 +10,9 @@ The build ID can be found in `Core\Constants.cs` in the `BuildId` variable - it
 
 The build ID moniker is made up of two components - a reference to a [Halo](https://en.wikipedia.org/wiki/Halo_(franchise)) character, and the date when the work on the specific build started in the format of `MMDDYYYY`.
 
-| Build ID                                                           | Build Date         |
+| Build ID                                                           | Build Date        |
 |:-------------------------------------------------------------------|:------------------|
+| [`TILLSON_11272024`](#TILLSON_11272024-november-27-2024)           | November 27, 2024 |
 | [`PROMETHEAN_09082024`](#PROMETHEAN_09082024-september-8-2024)     | September 8, 2024 |
 | [`VISEGRADRELAY_08152024`](#VISEGRADRELAY_08152024-august-15-2024) | August 15, 2024   |
 | [`DAISY023_04102024`](#DAISY023_04102024-april-10-2024)            | April 10, 2024    |
@@ -19,13 +20,28 @@ The build ID moniker is made up of two components - a reference to a [Halo](http
 | [`LIBRARIAN_03202022`](#librarian_03202022-march-20-2022)          | March 20, 2022    |
 | `ARBITER_01312022`                                                 | January 31, 2022  |
 
+### `TILLSON_11272024` (November 27, 2024)
+
+>[!NOTE]
+>See pull request: [Awake - `TILLSON_11272024`](https://github.com/microsoft/PowerToys/pull/36049)
+
+- [#35250](https://github.com/microsoft/PowerToys/issues/35250) Updates the icon retry policy, making sure that the icon consistently and correctly renders in the tray.
+- [#35848](https://github.com/microsoft/PowerToys/issues/35848) Fixed a bug where custom tray time shortcuts for longer than 24 hours would be parsed as zero hours/zero minutes.
+- [#34716](https://github.com/microsoft/PowerToys/issues/34716) Properly recover the state icon in the tray after an `explorer.exe` crash.
+- Added configuration safeguards to make sure that invalid values for timed keep-awake times do not result in exceptions.
+- Updated the tray initialization logic, making sure we wait for it to be properly created before setting icons.
+- Expanded logging capabilities to track invoking functions.
+- Added command validation logic to make sure that incorrect command line arguments display an error.
+- Display state now shown in the tray tooltip.
+- When timed mode is used, changing the display setting will no longer reset the timer.
+
 ### `PROMETHEAN_09082024` (September 8, 2024)
 
 >[!NOTE]
 >See pull request: [Awake - `PROMETHEAN_09082024`](https://github.com/microsoft/PowerToys/pull/34717)
 
 - Updating the initialization logic to make sure that settings are respected for proper group policy and single-instance detection.
-- [#34148] Fixed a bug from the previous release that incorrectly synchronized threads for shell icon creation and initialized parent PID when it was not parented.
+- [#34148](https://github.com/microsoft/PowerToys/issues/34148) Fixed a bug from the previous release that incorrectly synchronized threads for shell icon creation and initialized parent PID when it was not parented.
 
 ### `VISEGRADRELAY_08152024` (August 15, 2024)
 

src/modules/awake/Awake/Core/Constants.cs

@@ -17,6 +17,6 @@ namespace Awake.Core
         // Format of the build ID is: CODENAME_MMDDYYYY, where MMDDYYYY
         // is representative of the date when the last change was made before
         // the pull request is issued.
-        internal const string BuildId = "PROMETHEAN_09082024";
+        internal const string BuildId = "TILLSON_11272024";
     }
 }

src/modules/awake/Awake/Core/ExtensionMethods.cs

@@ -14,7 +14,7 @@ namespace Awake.Core
             ArgumentNullException.ThrowIfNull(target);
             ArgumentNullException.ThrowIfNull(source);
 
-            foreach (var element in source)
+            foreach (T? element in source)
             {
                 target.Add(element);
             }

src/modules/awake/Awake/Core/Manager.cs

@@ -10,11 +10,11 @@ using System.Drawing;
 using System.Globalization;
 using System.IO;
 using System.Reactive.Linq;
+using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 using System.Text;
 using System.Text.Json;
 using System.Threading;
-
 using Awake.Core.Models;
 using Awake.Core.Native;
 using Awake.Properties;
@@ -25,7 +25,7 @@ using Microsoft.Win32;
 
 namespace Awake.Core
 {
-    public delegate bool ConsoleEventHandler(Models.ControlType ctrlType);
+    public delegate bool ConsoleEventHandler(ControlType ctrlType);
 
     /// <summary>
     /// Helper class that allows talking to Win32 APIs without having to rely on PInvoke in other parts
@@ -33,27 +33,27 @@ namespace Awake.Core
     /// </summary>
     public class Manager
     {
-        private static bool _isUsingPowerToysConfig;
+        internal static bool IsUsingPowerToysConfig { get; set; }
 
-        internal static bool IsUsingPowerToysConfig { get => _isUsingPowerToysConfig; set => _isUsingPowerToysConfig = value; }
+        internal static SettingsUtils? ModuleSettings { get; set; }
+
+        private static AwakeMode CurrentOperatingMode { get; set; }
+
+        private static bool IsDisplayOn { get; set; }
+
+        private static uint TimeRemaining { get; set; }
+
+        private static string ScreenStateString => IsDisplayOn ? Resources.AWAKE_SCREEN_ON : Resources.AWAKE_SCREEN_OFF;
+
+        private static int ProcessId { get; set; }
+
+        private static DateTimeOffset ExpireAt { get; set; }
 
         private static readonly CompositeFormat AwakeMinutes = CompositeFormat.Parse(Resources.AWAKE_MINUTES);
         private static readonly CompositeFormat AwakeHours = CompositeFormat.Parse(Resources.AWAKE_HOURS);
-
         private static readonly BlockingCollection<ExecutionState> _stateQueue;
-
-        // Core icons used for the tray
-        private static readonly Icon _timedIcon = new(Path.Combine(AppDomain.CurrentDomain.BaseDirectory, "Assets/Awake/timed.ico"));
-        private static readonly Icon _expirableIcon = new(Path.Combine(AppDomain.CurrentDomain.BaseDirectory, "Assets/Awake/expirable.ico"));
-        private static readonly Icon _indefiniteIcon = new(Path.Combine(AppDomain.CurrentDomain.BaseDirectory, "Assets/Awake/indefinite.ico"));
-        private static readonly Icon _disabledIcon = new(Path.Combine(AppDomain.CurrentDomain.BaseDirectory, "Assets/Awake/disabled.ico"));
-
         private static CancellationTokenSource _tokenSource;
 
-        private static SettingsUtils? _moduleSettings;
-
-        internal static SettingsUtils? ModuleSettings { get => _moduleSettings; set => _moduleSettings = value; }
-
         static Manager()
         {
             _tokenSource = new CancellationTokenSource();
@@ -87,7 +87,7 @@ namespace Awake.Core
         {
             Bridge.AllocConsole();
 
-            var outputFilePointer = Bridge.CreateFile("CONOUT$", Native.Constants.GENERIC_READ | Native.Constants.GENERIC_WRITE, FileShare.Write, IntPtr.Zero, FileMode.OpenOrCreate, 0, IntPtr.Zero);
+            nint outputFilePointer = Bridge.CreateFile("CONOUT$", Native.Constants.GENERIC_READ | Native.Constants.GENERIC_WRITE, FileShare.Write, IntPtr.Zero, FileMode.OpenOrCreate, 0, IntPtr.Zero);
 
             Bridge.SetStdHandle(Native.Constants.STD_OUTPUT_HANDLE, outputFilePointer);
 
@@ -105,7 +105,7 @@ namespace Awake.Core
         {
             try
             {
-                var stateResult = Bridge.SetThreadExecutionState(state);
+                ExecutionState stateResult = Bridge.SetThreadExecutionState(state);
                 return stateResult != 0;
             }
             catch
@@ -123,42 +123,76 @@ namespace Awake.Core
 
         internal static void CancelExistingThread()
         {
-            Logger.LogInfo($"Attempting to ensure that the thread is properly cleaned up...");
+            Logger.LogInfo("Ensuring the thread is properly cleaned up...");
 
-            // Resetting the thread state.
+            // Reset the thread state and handle cancellation.
             _stateQueue.Add(ExecutionState.ES_CONTINUOUS);
 
-            // Next, make sure that any existing background threads are terminated.
             if (_tokenSource != null)
             {
                 _tokenSource.Cancel();
                 _tokenSource.Dispose();
-
-                _tokenSource = new CancellationTokenSource();
             }
             else
             {
-                Logger.LogWarning("The token source was null.");
+                Logger.LogWarning("Token source is null.");
             }
 
-            Logger.LogInfo("Instantiating of new token source and thread token completed.");
+            _tokenSource = new CancellationTokenSource();
+
+            Logger.LogInfo("New token source and thread token instantiated.");
         }
 
-        internal static void SetIndefiniteKeepAwake(bool keepDisplayOn = false)
+        internal static void SetModeShellIcon(bool forceAdd = false)
+        {
+            string iconText = string.Empty;
+            Icon? icon = null;
+
+            switch (CurrentOperatingMode)
+            {
+                case AwakeMode.INDEFINITE:
+                    string processText = ProcessId == 0
+                        ? string.Empty
+                        : $" - {Resources.AWAKE_TRAY_TEXT_PID_BINDING}: {ProcessId}";
+                    iconText = $"{Constants.FullAppName} [{Resources.AWAKE_TRAY_TEXT_INDEFINITE}{processText}][{ScreenStateString}]";
+                    icon = TrayHelper.IndefiniteIcon;
+                    break;
+
+                case AwakeMode.PASSIVE:
+                    iconText = $"{Constants.FullAppName} [{Resources.AWAKE_TRAY_TEXT_OFF}]";
+                    icon = TrayHelper.DisabledIcon;
+                    break;
+
+                case AwakeMode.EXPIRABLE:
+                    iconText = $"{Constants.FullAppName} [{Resources.AWAKE_TRAY_TEXT_EXPIRATION}][{ScreenStateString}][{ExpireAt:yyyy-MM-dd HH:mm:ss}]";
+                    icon = TrayHelper.ExpirableIcon;
+                    break;
+
+                case AwakeMode.TIMED:
+                    iconText = $"{Constants.FullAppName} [{Resources.AWAKE_TRAY_TEXT_TIMED}][{ScreenStateString}]";
+                    icon = TrayHelper.TimedIcon;
+                    break;
+            }
+
+            TrayHelper.SetShellIcon(
+                TrayHelper.WindowHandle,
+                iconText,
+                icon,
+                forceAdd ? TrayIconAction.Add : TrayIconAction.Update);
+        }
+
+        internal static void SetIndefiniteKeepAwake(bool keepDisplayOn = false, int processId = 0, [CallerMemberName] string callerName = "")
         {
             PowerToysTelemetry.Log.WriteEvent(new Telemetry.AwakeIndefinitelyKeepAwakeEvent());
 
-            TrayHelper.SetShellIcon(TrayHelper.HiddenWindowHandle, $"{Constants.FullAppName} [{Resources.AWAKE_TRAY_TEXT_INDEFINITE}]", _indefiniteIcon, TrayIconAction.Update);
-
             CancelExistingThread();
-            _stateQueue.Add(ComputeAwakeState(keepDisplayOn));
 
             if (IsUsingPowerToysConfig)
             {
                 try
                 {
-                    var currentSettings = ModuleSettings!.GetSettings<AwakeSettings>(Constants.AppName) ?? new AwakeSettings();
+                    AwakeSettings currentSettings = ModuleSettings!.GetSettings<AwakeSettings>(Constants.AppName) ?? new AwakeSettings();
-                    var settingsChanged = currentSettings.Properties.Mode != AwakeMode.INDEFINITE ||
+                    bool settingsChanged = currentSettings.Properties.Mode != AwakeMode.INDEFINITE ||
                                           currentSettings.Properties.KeepDisplayOn != keepDisplayOn;
 
                     if (settingsChanged)
@@ -166,6 +200,57 @@ namespace Awake.Core
                         currentSettings.Properties.Mode = AwakeMode.INDEFINITE;
                         currentSettings.Properties.KeepDisplayOn = keepDisplayOn;
                         ModuleSettings!.SaveSettings(JsonSerializer.Serialize(currentSettings), Constants.AppName);
+
+                        // We return here because when the settings are saved, they will be automatically
+                        // processed. That means that when they are processed, the indefinite keep-awake will kick-in properly
+                        // and we avoid double execution.
+                        return;
+                    }
+                }
+                catch (Exception ex)
+                {
+                    Logger.LogError($"Failed to handle indefinite keep awake command invoked by {callerName}: {ex.Message}");
+                }
+            }
+
+            Logger.LogInfo($"Indefinite keep-awake starting, invoked by {callerName}...");
+
+            _stateQueue.Add(ComputeAwakeState(keepDisplayOn));
+
+            IsDisplayOn = keepDisplayOn;
+            CurrentOperatingMode = AwakeMode.INDEFINITE;
+            ProcessId = processId;
+
+            SetModeShellIcon();
+        }
+
+        internal static void SetExpirableKeepAwake(DateTimeOffset expireAt, bool keepDisplayOn = true, [CallerMemberName] string callerName = "")
+        {
+            Logger.LogInfo($"Expirable keep-awake invoked by {callerName}. Expected expiration date/time: {expireAt} with display on setting set to {keepDisplayOn}.");
+            PowerToysTelemetry.Log.WriteEvent(new Telemetry.AwakeExpirableKeepAwakeEvent());
+
+            CancelExistingThread();
+
+            if (IsUsingPowerToysConfig)
+            {
+                try
+                {
+                    AwakeSettings currentSettings = ModuleSettings!.GetSettings<AwakeSettings>(Constants.AppName) ?? new AwakeSettings();
+                    bool settingsChanged = currentSettings.Properties.Mode != AwakeMode.EXPIRABLE ||
+                                          currentSettings.Properties.ExpirationDateTime != expireAt ||
+                                          currentSettings.Properties.KeepDisplayOn != keepDisplayOn;
+
+                    if (settingsChanged)
+                    {
+                        currentSettings.Properties.Mode = AwakeMode.EXPIRABLE;
+                        currentSettings.Properties.KeepDisplayOn = keepDisplayOn;
+                        currentSettings.Properties.ExpirationDateTime = expireAt;
+                        ModuleSettings!.SaveSettings(JsonSerializer.Serialize(currentSettings), Constants.AppName);
+
+                        // We return here because when the settings are saved, they will be automatically
+                        // processed. That means that when they are processed, the expirable keep-awake will kick-in properly
+                        // and we avoid double execution.
+                        return;
                     }
                 }
                 catch (Exception ex)
@@ -173,27 +258,30 @@ namespace Awake.Core
                     Logger.LogError($"Failed to handle indefinite keep awake command: {ex.Message}");
                 }
             }
-        }
 
-        internal static void SetExpirableKeepAwake(DateTimeOffset expireAt, bool keepDisplayOn = true)
+            Logger.LogInfo($"Expirable keep-awake starting...");
-        {
-            Logger.LogInfo($"Expirable keep-awake. Expected expiration date/time: {expireAt} with display on setting set to {keepDisplayOn}.");
 
-            PowerToysTelemetry.Log.WriteEvent(new Telemetry.AwakeExpirableKeepAwakeEvent());
+            if (expireAt <= DateTimeOffset.Now)
-
-            CancelExistingThread();
-
-            if (expireAt > DateTimeOffset.Now)
             {
-                Logger.LogInfo($"Starting expirable log for {expireAt}");
+                Logger.LogError($"The specified target date and time is not in the future. Current time: {DateTimeOffset.Now}, Target time: {expireAt}");
-                _stateQueue.Add(ComputeAwakeState(keepDisplayOn));
+                return;
+            }
 
-                TrayHelper.SetShellIcon(TrayHelper.HiddenWindowHandle, $"{Constants.FullAppName} [{Resources.AWAKE_TRAY_TEXT_EXPIRATION} - {expireAt}]", _expirableIcon, TrayIconAction.Update);
+            Logger.LogInfo($"Starting expirable log for {expireAt}");
+            _stateQueue.Add(ComputeAwakeState(keepDisplayOn));
 
-                Observable.Timer(expireAt - DateTimeOffset.Now).Subscribe(
+            IsDisplayOn = keepDisplayOn;
+            CurrentOperatingMode = AwakeMode.EXPIRABLE;
+            ExpireAt = expireAt;
+
+            SetModeShellIcon();
+
+            TimeSpan remainingTime = expireAt - DateTimeOffset.Now;
+
+            Observable.Timer(remainingTime).Subscribe(
                 _ =>
                 {
-                    Logger.LogInfo($"Completed expirable keep-awake.");
+                    Logger.LogInfo("Completed expirable keep-awake.");
                     CancelExistingThread();
 
                     if (IsUsingPowerToysConfig)
@@ -207,67 +295,85 @@ namespace Awake.Core
                     }
                 },
                 _tokenSource.Token);
-            }
+        }
-            else
+
-            {
+        internal static void SetTimedKeepAwake(uint seconds, bool keepDisplayOn = true, [CallerMemberName] string callerName = "")
-                Logger.LogError("The specified target date and time is not in the future.");
+        {
-                Logger.LogError($"Current time: {DateTimeOffset.Now}\tTarget time: {expireAt}");
+            Logger.LogInfo($"Timed keep-awake invoked by {callerName}. Expected runtime: {seconds} seconds with display on setting set to {keepDisplayOn}.");
-            }
+            PowerToysTelemetry.Log.WriteEvent(new Telemetry.AwakeTimedKeepAwakeEvent());
+
+            CancelExistingThread();
 
             if (IsUsingPowerToysConfig)
             {
                 try
                 {
-                    var currentSettings = ModuleSettings!.GetSettings<AwakeSettings>(Constants.AppName) ?? new AwakeSettings();
+                    AwakeSettings currentSettings = ModuleSettings!.GetSettings<AwakeSettings>(Constants.AppName) ?? new AwakeSettings();
-                    var settingsChanged = currentSettings.Properties.Mode != AwakeMode.EXPIRABLE ||
+                    TimeSpan timeSpan = TimeSpan.FromSeconds(seconds);
-                                          currentSettings.Properties.ExpirationDateTime != expireAt ||
+
-                                          currentSettings.Properties.KeepDisplayOn != keepDisplayOn;
+                    uint totalHours = (uint)timeSpan.TotalHours;
+                    uint remainingMinutes = (uint)Math.Ceiling(timeSpan.TotalMinutes % 60);
+
+                    bool settingsChanged = currentSettings.Properties.Mode != AwakeMode.TIMED ||
+                                          currentSettings.Properties.IntervalHours != totalHours ||
+                                          currentSettings.Properties.IntervalMinutes != remainingMinutes;
 
                     if (settingsChanged)
                     {
-                        currentSettings.Properties.Mode = AwakeMode.EXPIRABLE;
+                        currentSettings.Properties.Mode = AwakeMode.TIMED;
-                        currentSettings.Properties.KeepDisplayOn = keepDisplayOn;
+                        currentSettings.Properties.IntervalHours = totalHours;
-                        currentSettings.Properties.ExpirationDateTime = expireAt;
+                        currentSettings.Properties.IntervalMinutes = remainingMinutes;
                         ModuleSettings!.SaveSettings(JsonSerializer.Serialize(currentSettings), Constants.AppName);
+
+                        // We return here because when the settings are saved, they will be automatically
+                        // processed. That means that when they are processed, the timed keep-awake will kick-in properly
+                        // and we avoid double execution.
+                        return;
                     }
                 }
                 catch (Exception ex)
                 {
-                    Logger.LogError($"Failed to handle indefinite keep awake command: {ex.Message}");
+                    Logger.LogError($"Failed to handle timed keep awake command: {ex.Message}");
                 }
             }
-        }
 
-        internal static void SetTimedKeepAwake(uint seconds, bool keepDisplayOn = true)
+            Logger.LogInfo($"Timed keep-awake starting...");
-        {
-            Logger.LogInfo($"Timed keep-awake. Expected runtime: {seconds} seconds with display on setting set to {keepDisplayOn}.");
 
-            PowerToysTelemetry.Log.WriteEvent(new Telemetry.AwakeTimedKeepAwakeEvent());
-
-            CancelExistingThread();
-
-            Logger.LogInfo($"Timed keep awake started for {seconds} seconds.");
             _stateQueue.Add(ComputeAwakeState(keepDisplayOn));
 
-            TrayHelper.SetShellIcon(TrayHelper.HiddenWindowHandle, $"{Constants.FullAppName} [{Resources.AWAKE_TRAY_TEXT_TIMED}]", _timedIcon, TrayIconAction.Update);
+            IsDisplayOn = keepDisplayOn;
+            CurrentOperatingMode = AwakeMode.TIMED;
 
-            var timerObservable = Observable.Timer(TimeSpan.FromSeconds(seconds));
+            SetModeShellIcon();
-            var intervalObservable = Observable.Interval(TimeSpan.FromSeconds(1)).TakeUntil(timerObservable);
 
-            var combinedObservable = Observable.CombineLatest(intervalObservable, timerObservable.StartWith(0), (elapsedSeconds, _) => elapsedSeconds + 1);
+            ulong desiredDuration = (ulong)seconds * 1000;
+            ulong targetDuration = Math.Min(desiredDuration, uint.MaxValue - 1) / 1000;
+
+            if (desiredDuration > uint.MaxValue)
+            {
+                Logger.LogInfo($"The desired interval of {seconds} seconds ({desiredDuration}ms) exceeds the limit. Defaulting to maximum possible value: {targetDuration} seconds. Read more about existing limits in the official documentation: https://aka.ms/powertoys/awake");
+            }
+
+            IObservable<long> timerObservable = Observable.Timer(TimeSpan.FromSeconds(targetDuration));
+            IObservable<long> intervalObservable = Observable.Interval(TimeSpan.FromSeconds(1)).TakeUntil(timerObservable);
+            IObservable<long> combinedObservable = Observable.CombineLatest(intervalObservable, timerObservable.StartWith(0), (elapsedSeconds, _) => elapsedSeconds + 1);
 
             combinedObservable.Subscribe(
                 elapsedSeconds =>
                 {
-                    var timeRemaining = seconds - (uint)elapsedSeconds;
+                    TimeRemaining = (uint)targetDuration - (uint)elapsedSeconds;
-                    if (timeRemaining >= 0)
+                    if (TimeRemaining >= 0)
                     {
-                        TrayHelper.SetShellIcon(TrayHelper.HiddenWindowHandle, $"{Constants.FullAppName} [{Resources.AWAKE_TRAY_TEXT_TIMED}]\n{TimeSpan.FromSeconds(timeRemaining).ToHumanReadableString()}", _timedIcon, TrayIconAction.Update);
+                        TrayHelper.SetShellIcon(
+                            TrayHelper.WindowHandle,
+                            $"{Constants.FullAppName} [{Resources.AWAKE_TRAY_TEXT_TIMED}][{ScreenStateString}][{TimeSpan.FromSeconds(TimeRemaining).ToHumanReadableString()}]",
+                            TrayHelper.TimedIcon,
+                            TrayIconAction.Update);
                     }
                 },
                 () =>
                 {
-                    Console.WriteLine("Completed timed thread.");
+                    Logger.LogInfo("Completed timed thread.");
                     CancelExistingThread();
 
                     if (IsUsingPowerToysConfig)
@@ -283,30 +389,6 @@ namespace Awake.Core
                     }
                 },
                 _tokenSource.Token);
-
-            if (IsUsingPowerToysConfig)
-            {
-                try
-                {
-                    var currentSettings = ModuleSettings!.GetSettings<AwakeSettings>(Constants.AppName) ?? new AwakeSettings();
-                    var timeSpan = TimeSpan.FromSeconds(seconds);
-                    var settingsChanged = currentSettings.Properties.Mode != AwakeMode.TIMED ||
-                                          currentSettings.Properties.IntervalHours != (uint)timeSpan.Hours ||
-                                          currentSettings.Properties.IntervalMinutes != (uint)timeSpan.Minutes;
-
-                    if (settingsChanged)
-                    {
-                        currentSettings.Properties.Mode = AwakeMode.TIMED;
-                        currentSettings.Properties.IntervalHours = (uint)timeSpan.Hours;
-                        currentSettings.Properties.IntervalMinutes = (uint)timeSpan.Minutes;
-                        ModuleSettings!.SaveSettings(JsonSerializer.Serialize(currentSettings), Constants.AppName);
-                    }
-                }
-                catch (Exception ex)
-                {
-                    Logger.LogError($"Failed to handle timed keep awake command: {ex.Message}");
-                }
-            }
         }
 
         /// <summary>
@@ -317,15 +399,15 @@ namespace Awake.Core
         {
             SetPassiveKeepAwake(updateSettings: false);
 
-            if (TrayHelper.HiddenWindowHandle != IntPtr.Zero)
+            if (TrayHelper.WindowHandle != IntPtr.Zero)
             {
                 // Delete the icon.
-                TrayHelper.SetShellIcon(TrayHelper.HiddenWindowHandle, string.Empty, null, TrayIconAction.Delete);
+                TrayHelper.SetShellIcon(TrayHelper.WindowHandle, string.Empty, null, TrayIconAction.Delete);
 
                 // Close the message window that we used for the tray.
-                Bridge.SendMessage(TrayHelper.HiddenWindowHandle, Native.Constants.WM_CLOSE, 0, 0);
+                Bridge.SendMessage(TrayHelper.WindowHandle, Native.Constants.WM_CLOSE, 0, 0);
 
-                Bridge.DestroyWindow(TrayHelper.HiddenWindowHandle);
+                Bridge.DestroyWindow(TrayHelper.WindowHandle);
             }
 
             Bridge.PostQuitMessage(exitCode);
@@ -344,7 +426,7 @@ namespace Awake.Core
 
                 if (registryKey != null)
                 {
-                    var versionString = $"{registryKey.GetValue("ProductName")} {registryKey.GetValue("DisplayVersion")} {registryKey.GetValue("BuildLabEx")}";
+                    string versionString = $"{registryKey.GetValue("ProductName")} {registryKey.GetValue("DisplayVersion")} {registryKey.GetValue("BuildLabEx")}";
                     return versionString;
                 }
                 else
@@ -364,9 +446,9 @@ namespace Awake.Core
         /// Generates the default system tray options in situations where no custom options are provided.
         /// </summary>
         /// <returns>Returns a dictionary of default Awake timed interval options.</returns>
-        internal static Dictionary<string, int> GetDefaultTrayOptions()
+        internal static Dictionary<string, uint> GetDefaultTrayOptions()
         {
-            Dictionary<string, int> optionsList = new()
+            Dictionary<string, uint> optionsList = new()
             {
                 { string.Format(CultureInfo.InvariantCulture, AwakeMinutes, 30), 1800 },
                 { string.Format(CultureInfo.InvariantCulture, AwakeHours, 1), 3600 },
@@ -379,26 +461,28 @@ namespace Awake.Core
         /// Resets the computer to standard power settings.
         /// </summary>
         /// <param name="updateSettings">In certain cases, such as exits, we want to make sure that settings are not reset for the passive mode but rather retained based on previous execution. Default is to save settings, but otherwise it can be overridden.</param>
-        internal static void SetPassiveKeepAwake(bool updateSettings = true)
+        internal static void SetPassiveKeepAwake(bool updateSettings = true, [CallerMemberName] string callerName = "")
         {
-            Logger.LogInfo($"Operating in passive mode (computer's standard power plan). No custom keep awake settings enabled.");
+            Logger.LogInfo($"Operating in passive mode (computer's standard power plan). Invoked by {callerName}. No custom keep awake settings enabled.");
-
             PowerToysTelemetry.Log.WriteEvent(new Telemetry.AwakeNoKeepAwakeEvent());
 
             CancelExistingThread();
 
-            TrayHelper.SetShellIcon(TrayHelper.HiddenWindowHandle, $"{Constants.FullAppName} [{Resources.AWAKE_TRAY_TEXT_OFF}]", _disabledIcon, TrayIconAction.Update);
-
             if (IsUsingPowerToysConfig && updateSettings)
             {
                 try
                 {
-                    var currentSettings = ModuleSettings!.GetSettings<AwakeSettings>(Constants.AppName) ?? new AwakeSettings();
+                    AwakeSettings currentSettings = ModuleSettings!.GetSettings<AwakeSettings>(Constants.AppName) ?? new AwakeSettings();
 
                     if (currentSettings.Properties.Mode != AwakeMode.PASSIVE)
                     {
                         currentSettings.Properties.Mode = AwakeMode.PASSIVE;
                         ModuleSettings!.SaveSettings(JsonSerializer.Serialize(currentSettings), Constants.AppName);
+
+                        // We return here because when the settings are saved, they will be automatically
+                        // processed. That means that when they are processed, the passive keep-awake will kick-in properly
+                        // and we avoid double execution.
+                        return;
                     }
                 }
                 catch (Exception ex)
@@ -406,19 +490,38 @@ namespace Awake.Core
                     Logger.LogError($"Failed to reset Awake mode: {ex.Message}");
                 }
             }
+
+            Logger.LogInfo($"Passive keep-awake starting...");
+
+            CurrentOperatingMode = AwakeMode.PASSIVE;
+
+            SetModeShellIcon();
         }
 
         /// <summary>
         /// Sets the display settings.
         /// </summary>
-        internal static void SetDisplay()
+        internal static void SetDisplay([CallerMemberName] string callerName = "")
         {
+            Logger.LogInfo($"Setting display configuration from settings. Invoked by {callerName}.");
             if (IsUsingPowerToysConfig)
             {
                 try
                 {
-                    var currentSettings = ModuleSettings!.GetSettings<AwakeSettings>(Constants.AppName) ?? new AwakeSettings();
+                    AwakeSettings currentSettings = ModuleSettings!.GetSettings<AwakeSettings>(Constants.AppName) ?? new AwakeSettings();
                     currentSettings.Properties.KeepDisplayOn = !currentSettings.Properties.KeepDisplayOn;
+
+                    // We want to make sure that if the display setting changes (e.g., through the tray)
+                    // then we do not reset the counter from zero. Because the settings are only storing
+                    // hours and minutes, we round up the minutes value up when changes occur.
+                    if (CurrentOperatingMode == AwakeMode.TIMED && TimeRemaining > 0)
+                    {
+                        TimeSpan timeSpan = TimeSpan.FromSeconds(TimeRemaining);
+
+                        currentSettings.Properties.IntervalHours = (uint)timeSpan.TotalHours;
+                        currentSettings.Properties.IntervalMinutes = (uint)Math.Ceiling(timeSpan.TotalMinutes % 60);
+                    }
+
                     ModuleSettings!.SaveSettings(JsonSerializer.Serialize(currentSettings), Constants.AppName);
                 }
                 catch (Exception ex)

src/modules/awake/Awake/Core/Models/SystemPowerCapabilities.cs

@@ -3,68 +3,111 @@
 // See the LICENSE file in the project root for more information.
 
 using System.Runtime.InteropServices;
+using System.Text.Json.Serialization;
 
 namespace Awake.Core.Models
 {
+    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Ansi)]
     public struct SystemPowerCapabilities
     {
         [MarshalAs(UnmanagedType.U1)]
         public bool PowerButtonPresent;
+
         [MarshalAs(UnmanagedType.U1)]
         public bool SleepButtonPresent;
+
         [MarshalAs(UnmanagedType.U1)]
         public bool LidPresent;
+
         [MarshalAs(UnmanagedType.U1)]
         public bool SystemS1;
+
         [MarshalAs(UnmanagedType.U1)]
         public bool SystemS2;
+
         [MarshalAs(UnmanagedType.U1)]
         public bool SystemS3;
+
         [MarshalAs(UnmanagedType.U1)]
         public bool SystemS4;
+
         [MarshalAs(UnmanagedType.U1)]
         public bool SystemS5;
+
         [MarshalAs(UnmanagedType.U1)]
         public bool HiberFilePresent;
+
         [MarshalAs(UnmanagedType.U1)]
         public bool FullWake;
+
         [MarshalAs(UnmanagedType.U1)]
         public bool VideoDimPresent;
+
         [MarshalAs(UnmanagedType.U1)]
         public bool ApmPresent;
+
         [MarshalAs(UnmanagedType.U1)]
         public bool UpsPresent;
+
         [MarshalAs(UnmanagedType.U1)]
         public bool ThermalControl;
+
         [MarshalAs(UnmanagedType.U1)]
         public bool ProcessorThrottle;
+
         public byte ProcessorMinThrottle;
+        public byte ProcessorThrottleScale;
+
+        [MarshalAs(UnmanagedType.ByValArray, SizeConst = 4)]
+        public byte[] Spare2;
+
         public byte ProcessorMaxThrottle;
+
         [MarshalAs(UnmanagedType.U1)]
         public bool FastSystemS4;
+
         [MarshalAs(UnmanagedType.U1)]
         public bool Hiberboot;
+
         [MarshalAs(UnmanagedType.U1)]
         public bool WakeAlarmPresent;
+
         [MarshalAs(UnmanagedType.U1)]
         public bool AoAc;
+
         [MarshalAs(UnmanagedType.U1)]
         public bool DiskSpinDown;
+
+        [MarshalAs(UnmanagedType.ByValArray, SizeConst = 8)]
+        public byte[] Spare3;
+
         public byte HiberFileType;
+
         [MarshalAs(UnmanagedType.U1)]
         public bool AoAcConnectivitySupported;
-        [MarshalAs(UnmanagedType.ByValArray, SizeConst = 6)]
+
-        private readonly byte[] spare3;
         [MarshalAs(UnmanagedType.U1)]
         public bool SystemBatteriesPresent;
+
         [MarshalAs(UnmanagedType.U1)]
         public bool BatteriesAreShortTerm;
+
         [MarshalAs(UnmanagedType.ByValArray, SizeConst = 3)]
         public BatteryReportingScale[] BatteryScale;
+
+        [JsonConverter(typeof(JsonStringEnumConverter))]
         public SystemPowerState AcOnLineWake;
+
+        [JsonConverter(typeof(JsonStringEnumConverter))]
         public SystemPowerState SoftLidWake;
+
+        [JsonConverter(typeof(JsonStringEnumConverter))]
         public SystemPowerState RtcWake;
+
+        [JsonConverter(typeof(JsonStringEnumConverter))]
         public SystemPowerState MinDeviceWakeState;
+
+        [JsonConverter(typeof(JsonStringEnumConverter))]
         public SystemPowerState DefaultLowLatencyWake;
     }
 }

src/modules/awake/Awake/Core/Native/Bridge.cs

@@ -106,5 +106,8 @@ namespace Awake.Core.Native
 
         [DllImport("ntdll.dll")]
         internal static extern int NtQueryInformationProcess(IntPtr processHandle, int processInformationClass, ref ProcessBasicInformation processInformation, int processInformationLength, out int returnLength);
+
+        [DllImport("user32.dll", CharSet = CharSet.Unicode)]
+        internal static extern int RegisterWindowMessage(string lpString);
     }
 }

src/modules/awake/Awake/Core/Native/Constants.cs

@@ -11,6 +11,7 @@ namespace Awake.Core.Native
         internal const uint WM_COMMAND = 0x0111;
         internal const uint WM_USER = 0x0400U;
         internal const uint WM_CLOSE = 0x0010;
+        internal const int WM_CREATE = 0x0001;
         internal const int WM_DESTROY = 0x0002;
         internal const int WM_LBUTTONDOWN = 0x0201;
         internal const int WM_RBUTTONDOWN = 0x0204;

src/modules/awake/Awake/Core/Threading/SingleThreadSynchronizationContext.cs

@@ -5,7 +5,6 @@
 using System;
 using System.Collections.Generic;
 using System.Threading;
-
 using ManagedCommon;
 
 namespace Awake.Core.Threading

src/modules/awake/Awake/Core/TrayHelper.cs

@@ -6,10 +6,12 @@ using System;
 using System.Collections.Generic;
 using System.ComponentModel;
 using System.Drawing;
+using System.IO;
 using System.Linq;
+using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 using System.Threading;
-
+using System.Threading.Tasks;
 using Awake.Core.Models;
 using Awake.Core.Native;
 using Awake.Core.Threading;
@@ -29,19 +31,24 @@ namespace Awake.Core
     internal static class TrayHelper
     {
         private static NotifyIconData _notifyIconData;
-        private static IntPtr _trayMenu;
-        private static IntPtr _hiddenWindowHandle;
         private static SingleThreadSynchronizationContext? _syncContext;
         private static Thread? _mainThread;
+        private static uint _taskbarCreatedMessage;
 
-        private static IntPtr TrayMenu { get => _trayMenu; set => _trayMenu = value; }
+        private static IntPtr TrayMenu { get; set; }
 
-        internal static IntPtr HiddenWindowHandle { get => _hiddenWindowHandle; private set => _hiddenWindowHandle = value; }
+        internal static IntPtr WindowHandle { get; private set; }
+
+        internal static readonly Icon DefaultAwakeIcon = new(Path.Combine(AppDomain.CurrentDomain.BaseDirectory, "Assets/Awake/awake.ico"));
+        internal static readonly Icon TimedIcon = new(Path.Combine(AppDomain.CurrentDomain.BaseDirectory, "Assets/Awake/timed.ico"));
+        internal static readonly Icon ExpirableIcon = new(Path.Combine(AppDomain.CurrentDomain.BaseDirectory, "Assets/Awake/expirable.ico"));
+        internal static readonly Icon IndefiniteIcon = new(Path.Combine(AppDomain.CurrentDomain.BaseDirectory, "Assets/Awake/indefinite.ico"));
+        internal static readonly Icon DisabledIcon = new(Path.Combine(AppDomain.CurrentDomain.BaseDirectory, "Assets/Awake/disabled.ico"));
 
         static TrayHelper()
         {
             TrayMenu = IntPtr.Zero;
-            HiddenWindowHandle = IntPtr.Zero;
+            WindowHandle = IntPtr.Zero;
         }
 
         private static void ShowContextMenu(IntPtr hWnd)
@@ -59,7 +66,7 @@ namespace Awake.Core
             Bridge.ScreenToClient(hWnd, ref cursorPos);
 
             // Set menu information
-            var menuInfo = new MenuInfo
+            MenuInfo menuInfo = new()
             {
                 CbSize = (uint)Marshal.SizeOf<MenuInfo>(),
                 FMask = Native.Constants.MIM_STYLE,
@@ -77,8 +84,10 @@ namespace Awake.Core
                   IntPtr.Zero);
         }
 
-        public static void InitializeTray(Icon icon, string text)
+        public static Task InitializeTray(Icon icon, string text)
         {
+            TaskCompletionSource<bool> trayInitialized = new();
+
             IntPtr hWnd = IntPtr.Zero;
 
             // Start the message loop asynchronously
@@ -89,53 +98,63 @@ namespace Awake.Core
 
                 RunOnMainThread(() =>
                 {
-                    WndClassEx wcex = new()
+                    try
                     {
-                        CbSize = (uint)Marshal.SizeOf(typeof(WndClassEx)),
+                        WndClassEx wcex = new()
-                        Style = 0,
+                        {
-                        LpfnWndProc = Marshal.GetFunctionPointerForDelegate<Bridge.WndProcDelegate>(WndProc),
+                            CbSize = (uint)Marshal.SizeOf<WndClassEx>(),
-                        CbClsExtra = 0,
+                            Style = 0,
-                        CbWndExtra = 0,
+                            LpfnWndProc = Marshal.GetFunctionPointerForDelegate<Bridge.WndProcDelegate>(WndProc),
-                        HInstance = Marshal.GetHINSTANCE(typeof(Program).Module),
+                            CbClsExtra = 0,
-                        HIcon = IntPtr.Zero,
+                            CbWndExtra = 0,
-                        HCursor = IntPtr.Zero,
+                            HInstance = Marshal.GetHINSTANCE(typeof(Program).Module),
-                        HbrBackground = IntPtr.Zero,
+                            HIcon = IntPtr.Zero,
-                        LpszMenuName = string.Empty,
+                            HCursor = IntPtr.Zero,
-                        LpszClassName = Constants.TrayWindowId,
+                            HbrBackground = IntPtr.Zero,
-                        HIconSm = IntPtr.Zero,
+                            LpszMenuName = string.Empty,
-                    };
+                            LpszClassName = Constants.TrayWindowId,
+                            HIconSm = IntPtr.Zero,
+                        };
 
-                    Bridge.RegisterClassEx(ref wcex);
+                        Bridge.RegisterClassEx(ref wcex);
 
-                    hWnd = Bridge.CreateWindowEx(
+                        hWnd = Bridge.CreateWindowEx(
-                        0,
+                            0,
-                        Constants.TrayWindowId,
+                            Constants.TrayWindowId,
-                        text,
+                            text,
-                        0x00CF0000 | 0x00000001 | 0x00000008, // WS_OVERLAPPEDWINDOW | WS_VISIBLE | WS_MINIMIZEBOX
+                            0x00CF0000 | 0x00000001 | 0x00000008, // WS_OVERLAPPEDWINDOW | WS_VISIBLE | WS_MINIMIZEBOX
-                        0,
+                            0,
-                        0,
+                            0,
-                        0,
+                            0,
-                        0,
+                            0,
-                        unchecked(-3),
+                            IntPtr.Zero,
-                        IntPtr.Zero,
+                            IntPtr.Zero,
-                        Marshal.GetHINSTANCE(typeof(Program).Module),
+                            Marshal.GetHINSTANCE(typeof(Program).Module),
-                        IntPtr.Zero);
+                            IntPtr.Zero);
 
-                    if (hWnd == IntPtr.Zero)
+                        if (hWnd == IntPtr.Zero)
-                    {
+                        {
-                        int errorCode = Marshal.GetLastWin32Error();
+                            int errorCode = Marshal.GetLastWin32Error();
-                        throw new Win32Exception(errorCode, "Failed to add tray icon. Error code: " + errorCode);
+                            throw new Win32Exception(errorCode, "Failed to add tray icon. Error code: " + errorCode);
+                        }
+
+                        // Keep this as a reference because we will need it when we update
+                        // the tray icon in the future.
+                        WindowHandle = hWnd;
+
+                        Bridge.ShowWindow(hWnd, 0); // SW_HIDE
+                        Bridge.UpdateWindow(hWnd);
+                        Logger.LogInfo($"Created HWND for the window: {hWnd}");
+
+                        SetShellIcon(hWnd, text, icon);
+
+                        trayInitialized.SetResult(true);
+                    }
+                    catch (Exception ex)
+                    {
+                        Logger.LogError($"Failed to properly initialize the tray. {ex.Message}");
+                        trayInitialized.SetException(ex);
                     }
-
-                    // Keep this as a reference because we will need it when we update
-                    // the tray icon in the future.
-                    HiddenWindowHandle = hWnd;
-
-                    Bridge.ShowWindow(hWnd, 0); // SW_HIDE
-                    Bridge.UpdateWindow(hWnd);
-                    Logger.LogInfo($"Created HWND for the window: {hWnd}");
-
-                    SetShellIcon(hWnd, text, icon);
                 });
 
                 RunOnMainThread(() =>
@@ -148,9 +167,11 @@ namespace Awake.Core
 
             _mainThread.IsBackground = true;
             _mainThread.Start();
+
+            return trayInitialized.Task;
         }
 
-        internal static void SetShellIcon(IntPtr hWnd, string text, Icon? icon, TrayIconAction action = TrayIconAction.Add)
+        internal static void SetShellIcon(IntPtr hWnd, string text, Icon? icon, TrayIconAction action = TrayIconAction.Add, [CallerMemberName] string callerName = "")
         {
             if (hWnd != IntPtr.Zero && icon != null)
             {
@@ -169,11 +190,11 @@ namespace Awake.Core
                         break;
                 }
 
-                if (action == TrayIconAction.Add || action == TrayIconAction.Update)
+                if (action is TrayIconAction.Add or TrayIconAction.Update)
                 {
                     _notifyIconData = new NotifyIconData
                     {
-                        CbSize = Marshal.SizeOf(typeof(NotifyIconData)),
+                        CbSize = Marshal.SizeOf<NotifyIconData>(),
                         HWnd = hWnd,
                         UId = 1000,
                         UFlags = Native.Constants.NIF_ICON | Native.Constants.NIF_TIP | Native.Constants.NIF_MESSAGE,
@@ -186,19 +207,32 @@ namespace Awake.Core
                 {
                     _notifyIconData = new NotifyIconData
                     {
-                        CbSize = Marshal.SizeOf(typeof(NotifyIconData)),
+                        CbSize = Marshal.SizeOf<NotifyIconData>(),
                         HWnd = hWnd,
                         UId = 1000,
                         UFlags = 0,
                     };
                 }
 
-                if (!Bridge.Shell_NotifyIcon(message, ref _notifyIconData))
+                for (int attempt = 1; attempt <= 3; attempt++)
                 {
-                    int errorCode = Marshal.GetLastWin32Error();
+                    if (Bridge.Shell_NotifyIcon(message, ref _notifyIconData))
-                    Logger.LogInfo($"Could not set the shell icon. Action: {action} and error code: {errorCode}. HIcon handle is {icon?.Handle} and HWnd is {hWnd}");
+                    {
+                        break;
+                    }
+                    else
+                    {
+                        int errorCode = Marshal.GetLastWin32Error();
+                        Logger.LogInfo($"Could not set the shell icon. Action: {action}, error code: {errorCode}. HIcon handle is {icon?.Handle} and HWnd is {hWnd}. Invoked by {callerName}.");
 
-                    throw new Win32Exception(errorCode, $"Failed to change tray icon. Action: {action} and error code: {errorCode}");
+                        if (attempt == 3)
+                        {
+                            Logger.LogError($"Failed to change tray icon after 3 attempts. Action: {action} and error code: {errorCode}. Invoked by {callerName}.");
+                            break;
+                        }
+
+                        Thread.Sleep(100);
+                    }
                 }
 
                 if (action == TrayIconAction.Delete)
@@ -228,19 +262,26 @@ namespace Awake.Core
             switch (message)
             {
                 case Native.Constants.WM_USER:
-                    if (lParam == (IntPtr)Native.Constants.WM_LBUTTONDOWN || lParam == (IntPtr)Native.Constants.WM_RBUTTONDOWN)
+                    if (lParam is Native.Constants.WM_LBUTTONDOWN or Native.Constants.WM_RBUTTONDOWN)
                     {
                         // Show the context menu associated with the tray icon
                         ShowContextMenu(hWnd);
                     }
 
+                    break;
+
+                case Native.Constants.WM_CREATE:
+                    {
+                        _taskbarCreatedMessage = (uint)Bridge.RegisterWindowMessage("TaskbarCreated");
+                    }
+
                     break;
                 case Native.Constants.WM_DESTROY:
                     // Clean up resources when the window is destroyed
                     Bridge.PostQuitMessage(0);
                     break;
                 case Native.Constants.WM_COMMAND:
-                    int trayCommandsSize = Enum.GetNames(typeof(TrayCommands)).Length;
+                    int trayCommandsSize = Enum.GetNames<TrayCommands>().Length;
 
                     long targetCommandIndex = wParam.ToInt64() & 0xFFFF;
 
@@ -282,7 +323,7 @@ namespace Awake.Core
                                     }
 
                                     int index = (int)targetCommandIndex - (int)TrayCommands.TC_TIME;
-                                    uint targetTime = (uint)settings.Properties.CustomTrayTimes.ElementAt(index).Value;
+                                    uint targetTime = settings.Properties.CustomTrayTimes.ElementAt(index).Value;
                                     Manager.SetTimedKeepAwake(targetTime, keepDisplayOn: settings.Properties.KeepDisplayOn);
                                 }
 
@@ -292,6 +333,12 @@ namespace Awake.Core
 
                     break;
                 default:
+                    if (message == _taskbarCreatedMessage)
+                    {
+                        Logger.LogInfo("Taskbar re-created");
+                        Manager.SetModeShellIcon(forceAdd: true);
+                    }
+
                     // Let the default window procedure handle other messages
                     return Bridge.DefWindowProc(hWnd, message, wParam, lParam);
             }
@@ -326,7 +373,7 @@ namespace Awake.Core
                 startedFromPowerToys);
         }
 
-        public static void SetTray(bool keepDisplayOn, AwakeMode mode, Dictionary<string, int> trayTimeShortcuts, bool startedFromPowerToys)
+        public static void SetTray(bool keepDisplayOn, AwakeMode mode, Dictionary<string, uint> trayTimeShortcuts, bool startedFromPowerToys)
         {
             ClearExistingTrayMenu();
             CreateNewTrayMenu(startedFromPowerToys, keepDisplayOn, mode);
@@ -382,7 +429,7 @@ namespace Awake.Core
             Bridge.InsertMenu(TrayMenu, (uint)position, Native.Constants.MF_BYPOSITION | Native.Constants.MF_SEPARATOR, 0, string.Empty);
         }
 
-        private static void EnsureDefaultTrayTimeShortcuts(Dictionary<string, int> trayTimeShortcuts)
+        private static void EnsureDefaultTrayTimeShortcuts(Dictionary<string, uint> trayTimeShortcuts)
         {
             if (trayTimeShortcuts.Count == 0)
             {
@@ -390,15 +437,15 @@ namespace Awake.Core
             }
         }
 
-        private static void CreateAwakeTimeSubMenu(Dictionary<string, int> trayTimeShortcuts, bool isChecked = false)
+        private static void CreateAwakeTimeSubMenu(Dictionary<string, uint> trayTimeShortcuts, bool isChecked = false)
         {
-            var awakeTimeMenu = Bridge.CreatePopupMenu();
+            nint awakeTimeMenu = Bridge.CreatePopupMenu();
             for (int i = 0; i < trayTimeShortcuts.Count; i++)
             {
                 Bridge.InsertMenu(awakeTimeMenu, (uint)i, Native.Constants.MF_BYPOSITION | Native.Constants.MF_STRING, (uint)TrayCommands.TC_TIME + (uint)i, trayTimeShortcuts.ElementAt(i).Key);
             }
 
-            Bridge.InsertMenu(TrayMenu, 0, Native.Constants.MF_BYPOSITION | Native.Constants.MF_POPUP | (isChecked == true ? Native.Constants.MF_CHECKED : Native.Constants.MF_UNCHECKED), (uint)awakeTimeMenu, Resources.AWAKE_KEEP_ON_INTERVAL);
+            Bridge.InsertMenu(TrayMenu, 0, Native.Constants.MF_BYPOSITION | Native.Constants.MF_POPUP | (isChecked ? Native.Constants.MF_CHECKED : Native.Constants.MF_UNCHECKED), (uint)awakeTimeMenu, Resources.AWAKE_KEEP_ON_INTERVAL);
         }
 
         private static void InsertAwakeModeMenuItems(AwakeMode mode)

src/modules/awake/Awake/Program.cs

@@ -4,8 +4,8 @@
 
 using System;
 using System.CommandLine;
+using System.CommandLine.Parsing;
 using System.Diagnostics;
-using System.Drawing;
 using System.Globalization;
 using System.IO;
 using System.Linq;
@@ -15,7 +15,6 @@ using System.Reflection;
 using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
-
 using Awake.Core;
 using Awake.Core.Models;
 using Awake.Core.Native;
@@ -28,33 +27,34 @@ namespace Awake
 {
     internal sealed class Program
     {
-        private static Mutex? _mutex;
+        private static readonly string[] _aliasesConfigOption = ["--use-pt-config", "-c"];
+        private static readonly string[] _aliasesDisplayOption = ["--display-on", "-d"];
+        private static readonly string[] _aliasesTimeOption = ["--time-limit", "-t"];
+        private static readonly string[] _aliasesPidOption = ["--pid", "-p"];
+        private static readonly string[] _aliasesExpireAtOption = ["--expire-at", "-e"];
+        private static readonly string[] _aliasesParentPidOption = ["--use-parent-pid", "-u"];
+
+        private static readonly JsonSerializerOptions _serializerOptions = new() { IncludeFields = true };
+        private static readonly ETWTrace _etwTrace = new();
+
         private static FileSystemWatcher? _watcher;
         private static SettingsUtils? _settingsUtils;
-        private static ETWTrace _etwTrace = new ETWTrace();
 
         private static bool _startedFromPowerToys;
 
-        public static Mutex? LockMutex { get => _mutex; set => _mutex = value; }
+        public static Mutex? LockMutex { get; set; }
 
 #pragma warning disable CS8618 // Non-nullable field must contain a non-null value when exiting constructor. Consider declaring as nullable.
         private static ConsoleEventHandler _handler;
         private static SystemPowerCapabilities _powerCapabilities;
 #pragma warning restore CS8618 // Non-nullable field must contain a non-null value when exiting constructor. Consider declaring as nullable.
 
-        internal static readonly string[] AliasesConfigOption = ["--use-pt-config", "-c"];
+        private static async Task<int> Main(string[] args)
-        internal static readonly string[] AliasesDisplayOption = ["--display-on", "-d"];
-        internal static readonly string[] AliasesTimeOption = ["--time-limit", "-t"];
-        internal static readonly string[] AliasesPidOption = ["--pid", "-p"];
-        internal static readonly string[] AliasesExpireAtOption = ["--expire-at", "-e"];
-        internal static readonly string[] AliasesParentPidOption = ["--use-parent-pid", "-u"];
-
-        private static readonly Icon _defaultAwakeIcon = new(Path.Combine(AppDomain.CurrentDomain.BaseDirectory, "Assets/Awake/awake.ico"));
-
-        private static int Main(string[] args)
         {
             _settingsUtils = new SettingsUtils();
+
             LockMutex = new Mutex(true, Core.Constants.AppName, out bool instantiated);
+
             Logger.InitializeLogger(Path.Combine("\\", Core.Constants.AppName, "Logs"));
 
             try
@@ -62,7 +62,7 @@ namespace Awake
                 string appLanguage = LanguageHelper.LoadLanguage();
                 if (!string.IsNullOrEmpty(appLanguage))
                 {
-                    System.Threading.Thread.CurrentThread.CurrentUICulture = new CultureInfo(appLanguage);
+                    Thread.CurrentThread.CurrentUICulture = new CultureInfo(appLanguage);
                 }
             }
             catch (CultureNotFoundException ex)
@@ -70,6 +70,8 @@ namespace Awake
                 Logger.LogError("CultureNotFoundException: " + ex.Message);
             }
 
+            await TrayHelper.InitializeTray(TrayHelper.DefaultAwakeIcon, Core.Constants.FullAppName);
+            AppDomain.CurrentDomain.ProcessExit += (_, _) => TrayHelper.RunOnMainThread(() => LockMutex?.ReleaseMutex());
             AppDomain.CurrentDomain.UnhandledException += AwakeUnhandledExceptionCatcher;
 
             if (!instantiated)
@@ -103,46 +105,76 @@ namespace Awake
                     // To make it easier to diagnose future issues, let's get the
                     // system power capabilities and aggregate them in the log.
                     Bridge.GetPwrCapabilities(out _powerCapabilities);
-                    Logger.LogInfo(JsonSerializer.Serialize(_powerCapabilities));
+                    Logger.LogInfo(JsonSerializer.Serialize(_powerCapabilities, _serializerOptions));
 
                     Logger.LogInfo("Parsing parameters...");
 
-                    var configOption = new Option<bool>(AliasesConfigOption, () => false, Resources.AWAKE_CMD_HELP_CONFIG_OPTION)
+                    Option<bool> configOption = new(_aliasesConfigOption, () => false, Resources.AWAKE_CMD_HELP_CONFIG_OPTION)
                     {
                         Arity = ArgumentArity.ZeroOrOne,
                         IsRequired = false,
                     };
 
-                    var displayOption = new Option<bool>(AliasesDisplayOption, () => true, Resources.AWAKE_CMD_HELP_DISPLAY_OPTION)
+                    Option<bool> displayOption = new(_aliasesDisplayOption, () => true, Resources.AWAKE_CMD_HELP_DISPLAY_OPTION)
                     {
                         Arity = ArgumentArity.ZeroOrOne,
                         IsRequired = false,
                     };
 
-                    var timeOption = new Option<uint>(AliasesTimeOption, () => 0, Resources.AWAKE_CMD_HELP_TIME_OPTION)
+                    Option<uint> timeOption = new(_aliasesTimeOption, () => 0, Resources.AWAKE_CMD_HELP_TIME_OPTION)
                     {
                         Arity = ArgumentArity.ExactlyOne,
                         IsRequired = false,
                     };
 
-                    var pidOption = new Option<int>(AliasesPidOption, () => 0, Resources.AWAKE_CMD_HELP_PID_OPTION)
+                    Option<int> pidOption = new(_aliasesPidOption, () => 0, Resources.AWAKE_CMD_HELP_PID_OPTION)
                     {
                         Arity = ArgumentArity.ZeroOrOne,
                         IsRequired = false,
                     };
 
-                    var expireAtOption = new Option<string>(AliasesExpireAtOption, () => string.Empty, Resources.AWAKE_CMD_HELP_EXPIRE_AT_OPTION)
+                    Option<string> expireAtOption = new(_aliasesExpireAtOption, () => string.Empty, Resources.AWAKE_CMD_HELP_EXPIRE_AT_OPTION)
                     {
                         Arity = ArgumentArity.ZeroOrOne,
                         IsRequired = false,
                     };
 
-                    var parentPidOption = new Option<bool>(AliasesParentPidOption, () => false, Resources.AWAKE_CMD_PARENT_PID_OPTION)
+                    Option<bool> parentPidOption = new(_aliasesParentPidOption, () => false, Resources.AWAKE_CMD_PARENT_PID_OPTION)
                     {
                         Arity = ArgumentArity.ZeroOrOne,
                         IsRequired = false,
                     };
 
+                    timeOption.AddValidator(result =>
+                    {
+                        if (result.Tokens.Count != 0 && !uint.TryParse(result.Tokens[0].Value, out _))
+                        {
+                            string errorMessage = $"Interval in --time-limit could not be parsed correctly. Check that the value is valid and doesn't exceed 4,294,967,295. Value used: {result.Tokens[0].Value}.";
+                            Logger.LogError(errorMessage);
+                            result.ErrorMessage = errorMessage;
+                        }
+                    });
+
+                    pidOption.AddValidator(result =>
+                    {
+                        if (result.Tokens.Count != 0 && !int.TryParse(result.Tokens[0].Value, out _))
+                        {
+                            string errorMessage = $"PID value in --pid could not be parsed correctly. Check that the value is valid and falls within the boundaries of Windows PID process limits. Value used: {result.Tokens[0].Value}.";
+                            Logger.LogError(errorMessage);
+                            result.ErrorMessage = errorMessage;
+                        }
+                    });
+
+                    expireAtOption.AddValidator(result =>
+                    {
+                        if (result.Tokens.Count != 0 && !DateTimeOffset.TryParse(result.Tokens[0].Value, out _))
+                        {
+                            string errorMessage = $"Date and time value in --expire-at could not be parsed correctly. Check that the value is valid date and time. Refer to https://aka.ms/powertoys/awake for format examples. Value used: {result.Tokens[0].Value}.";
+                            Logger.LogError(errorMessage);
+                            result.ErrorMessage = errorMessage;
+                        }
+                    });
+
                     RootCommand? rootCommand =
                     [
                         configOption,
@@ -207,9 +239,7 @@ namespace Awake
             // Start the monitor thread that will be used to track the current state.
             Manager.StartMonitor();
 
-            TrayHelper.InitializeTray(_defaultAwakeIcon, Core.Constants.FullAppName);
+            EventWaitHandle eventHandle = new(false, EventResetMode.ManualReset, PowerToys.Interop.Constants.AwakeExitEvent());
-
-            var eventHandle = new EventWaitHandle(false, EventResetMode.ManualReset, PowerToys.Interop.Constants.AwakeExitEvent());
             new Thread(() =>
             {
                 WaitHandle.WaitAny([eventHandle]);
@@ -219,7 +249,8 @@ namespace Awake
             if (usePtConfig)
             {
                 // Configuration file is used, therefore we disregard any other command-line parameter
-                // and instead watch for changes in the file.
+                // and instead watch for changes in the file. This is used as a priority against all other arguments,
+                // so if --use-pt-config is applied the rest of the arguments are irrelevant.
                 Manager.IsUsingPowerToysConfig = true;
 
                 try
@@ -259,13 +290,14 @@ namespace Awake
                 // Second, we snap to process-based execution. Because this is something that
                 // is snapped to a running entity, we only want to enable the ability to set
                 // indefinite keep-awake with the display settings that the user wants to set.
+                // In this context, manual (explicit) PID takes precedence over parent PID.
                 int targetPid = pid != 0 ? pid : useParentPid ? Manager.GetParentProcess()?.Id ?? 0 : 0;
 
                 if (targetPid != 0)
                 {
                     Logger.LogInfo($"Bound to target process: {targetPid}");
 
-                    Manager.SetIndefiniteKeepAwake(displayOn);
+                    Manager.SetIndefiniteKeepAwake(displayOn, targetPid);
 
                     RunnerHelper.WaitForPowerToysRunner(targetPid, () =>
                     {
@@ -338,8 +370,8 @@ namespace Awake
 
         private static void SetupFileSystemWatcher(string settingsPath)
         {
-            var directory = Path.GetDirectoryName(settingsPath)!;
+            string directory = Path.GetDirectoryName(settingsPath)!;
-            var fileName = Path.GetFileName(settingsPath);
+            string fileName = Path.GetFileName(settingsPath);
 
             _watcher = new FileSystemWatcher
             {
@@ -364,7 +396,7 @@ namespace Awake
 
         private static void InitializeSettings()
         {
-            var settings = Manager.ModuleSettings?.GetSettings<AwakeSettings>(Core.Constants.AppName) ?? new AwakeSettings();
+            AwakeSettings settings = Manager.ModuleSettings?.GetSettings<AwakeSettings>(Core.Constants.AppName) ?? new AwakeSettings();
             TrayHelper.SetTray(settings, _startedFromPowerToys);
         }
 
@@ -385,7 +417,7 @@ namespace Awake
         {
             try
             {
-                var settings = _settingsUtils!.GetSettings<AwakeSettings>(Core.Constants.AppName)
+                AwakeSettings settings = _settingsUtils!.GetSettings<AwakeSettings>(Core.Constants.AppName)
                     ?? throw new InvalidOperationException("Settings are null.");
 
                 Logger.LogInfo($"Identified custom time shortcuts for the tray: {settings.Properties.CustomTrayTimes.Count}");

src/modules/awake/Awake/Properties/Resources.Designer.cs

@@ -258,6 +258,24 @@ namespace Awake.Properties {
             }
         }
         
+        /// <summary>
+        ///   Looks up a localized string similar to Off.
+        /// </summary>
+        internal static string AWAKE_SCREEN_OFF {
+            get {
+                return ResourceManager.GetString("AWAKE_SCREEN_OFF", resourceCulture);
+            }
+        }
+        
+        /// <summary>
+        ///   Looks up a localized string similar to On.
+        /// </summary>
+        internal static string AWAKE_SCREEN_ON {
+            get {
+                return ResourceManager.GetString("AWAKE_SCREEN_ON", resourceCulture);
+            }
+        }
+        
         /// <summary>
         ///   Looks up a localized string similar to Expiring.
         /// </summary>
@@ -285,6 +303,15 @@ namespace Awake.Properties {
             }
         }
         
+        /// <summary>
+        ///   Looks up a localized string similar to Bound to.
+        /// </summary>
+        internal static string AWAKE_TRAY_TEXT_PID_BINDING {
+            get {
+                return ResourceManager.GetString("AWAKE_TRAY_TEXT_PID_BINDING", resourceCulture);
+            }
+        }
+        
         /// <summary>
         ///   Looks up a localized string similar to Interval.
         /// </summary>

src/modules/awake/Awake/Properties/Resources.resx

@@ -208,4 +208,14 @@
   <data name="AWAKE_CMD_PARENT_PID_OPTION" xml:space="preserve">
     <value>Uses the parent process as the bound target - once the process terminates, Awake stops.</value>
   </data>
+  <data name="AWAKE_TRAY_TEXT_PID_BINDING" xml:space="preserve">
+    <value>Bound to</value>
+    <comment>Describes the process ID Awake is bound to when running.</comment>
+  </data>
+  <data name="AWAKE_SCREEN_ON" xml:space="preserve">
+    <value>On</value>
+  </data>
+  <data name="AWAKE_SCREEN_OFF" xml:space="preserve">
+    <value>Off</value>
+  </data>
 </root>

src/settings-ui/Settings.UI.Library/AwakeProperties.cs

@@ -38,7 +38,7 @@ namespace Microsoft.PowerToys.Settings.UI.Library
         public DateTimeOffset ExpirationDateTime { get; set; }
 
         [JsonPropertyName("customTrayTimes")]
-        [CmdConfigureIgnoreAttribute]
+        [CmdConfigureIgnore]
-        public Dictionary<string, int> CustomTrayTimes { get; set; }
+        public Dictionary<string, uint> CustomTrayTimes { get; set; }
     }
 }