[Awake] Context menu bug fixes (#16903)

* Fix the path to the icon

* Need the reverse when not working in isolation

* Updating some tray behaviors

* Making sure we have constants separately, and a filter

* Update tray logic

* Remove unnecessary menus

* Cleaning up how exit is done.

* Adding handling for tray commands

* Update with settings for dynamic times

* Proper reaction to timed keep-awake from the tray

* Proper handling for timed keep-awake from the tray

* Making sure that code analysis works correctly

* Making sure that errors are set in native calls

* Making sure the right icon path is used after testing

* Proper disposal of the context menu

* Fix tray designation

* Update with latest information on changes to the builds

* Update with guidance on files

* Update changelog doc

* Fix project file

* Remove `VTABLE`

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

@@ -193,6 +193,7 @@ buildtransitive
 Burkina
 Buryatia
 BValue
+BYPOSITION
 bytearray
 Caiguna
 CALG
@@ -212,6 +213,7 @@ cdpx
 CENTERALIGN
 cguid
 changecursor
+Changelog
 Changemove
 charconv
 chdir
@@ -374,7 +376,6 @@ DARKYELLOW
 datareader
 Datavalue
 DATAW
-David
 davidegiacometti
 Dayof
 Dbg
@@ -825,6 +826,7 @@ IFile
 IFilter
 ifndef
 IFolder
+IFormat
 ifstream
 IGraph
 iid
@@ -848,6 +850,7 @@ IMarkdown
 ime
 IMedia
 IMem
+IMessage
 imeutil
 iminstall
 IMoniker
@@ -1109,6 +1112,7 @@ LPSTR
 lpsz
 lpt
 LPTOP
+lptpm
 LPTSTR
 LPVOID
 LPW
@@ -1155,9 +1159,6 @@ MAXSHORTCUTSIZE
 maxversiontested
 Mbits
 MBs
-mdtext
-mdtxt
-mdwn
 MBUTTON
 MBUTTONDBLCLK
 MBUTTONDOWN
@@ -1167,12 +1168,16 @@ MCST
 MDICHILD
 MDL
 mdpreviewhandler
+mdtext
+mdtxt
+mdwn
 MEDIASUBTYPE
 mediatype
 Melman
 memcmp
 memcpy
 memset
+MENUBREAK
 MENUITEMINFO
 MENUITEMINFOW
 messageboxes
@@ -1223,7 +1228,6 @@ MONITORINFOEXW
 monitorinfof
 Monthand
 Moq
-Morton
 MOUSEACTIVATE
 MOUSEHWHEEL
 MOUSEINPUT
@@ -1381,6 +1385,7 @@ ntdll
 NTFS
 NTSTATUS
 nuget
+nuint
 nullopt
 nullptr
 numberbox
@@ -1439,6 +1444,7 @@ overlaywindow
 Overridable
 Oversampling
 OWNDC
+OWNERDRAW
 PACL
 pagos
 PAINTSTRUCT
@@ -1461,7 +1467,6 @@ pcb
 pch
 PCIDLIST
 PCWSTR
-
 pdb
 pdbonly
 pdfpreviewhandler
@@ -1731,6 +1736,7 @@ safeprojectname
 SAMEKEYPREVIOUSLYMAPPED
 SAMESHORTCUTPREVIOUSLYMAPPED
 SAVEFAILED
+scalability
 scancode
 scanled
 schedtasks
@@ -1966,8 +1972,8 @@ SYSKEYUP
 SYSLIB
 syslog
 SYSMENU
-systemd
 SYSTEMAPPS
+systemd
 SYSTEMTIME
 Tadele
 Tajikistan
@@ -2027,7 +2033,6 @@ Toolchain
 toolkitcontrols
 toolkitconverters
 Toolset
-toolstrip
 toolwindow
 TOPDOWNDIB
 toplevel
@@ -2187,7 +2192,6 @@ vstemplate
 VSTHRD
 VSTT
 vtable
-VTABLE
 Vtbl
 WBounds
 wca

doc/planning/awake.md

@@ -0,0 +1,38 @@
+---
+last-update: 3-20-2022
+---
+
+# PowerToys Awake Changelog
+
+## Builds
+
+The build ID can be found in [`NLog.config`](https://github.com/microsoft/PowerToys/blob/2e3a2b3f96f67c7dfc72963e5135662d3230b5fe/src/modules/awake/Awake/NLog.config#L5) - it is a unique identifier for the current builds that allows better diagnostics (we can look up the build ID from the logs) and offers a way to triage Awake-specific issues faster independent of the PowerToys version. The build ID does not carry any significance beyond that within the PowerToys code base.
+
+| Build ID                                                  | Build Date       |
+|:----------------------------------------------------------|:-----------------|
+| [`LIBRARIAN_03202022`](#librarian_03202022-march-20-2022) | March 20, 2022   |
+| `ARBITER_01312022`                                        | January 31, 2022 |
+
+### `LIBRARIAN_03202022` (March 20, 2022)
+
+- Changed the tray context menu to be following OS conventions instead of the style offered by Windows Forms. This introduces better support for DPI scaling and theming in the future.
+- Custom times in the tray can now be configured in the `settings.json` file for awake, through the `tray_times` property. The property values are representative of a `Dictionary<string, int>` and can be in the form of `"YOUR_NAME": LENGTH_IN_SECONDS`:
+
+```json
+{
+    "properties": {
+        "awake_keep_display_on": true,
+        "awake_mode": 2,
+        "awake_hours": 0,
+        "awake_minutes": 3,
+        "tray_times": {
+            "Custom length": 1800,
+            "Another custom length": 3600
+        }
+    },
+    "name": "Awake",
+    "version": "1.0"
+}
+```
+
+- Proper Awake background window closure was implemented to ensure that the process collects the correct handle instead of the empty one that was previously done through `System.Diagnostics.Process.GetCurrentProcess().CloseMainWindow()`. This likely can help with the Awake process that is left hanging after PowerToys itself closes.

src/modules/awake/Awake/Awake.csproj

@@ -2,7 +2,7 @@
   <Import Project="..\..\..\Version.props" />
   <PropertyGroup>
     <OutputType>WinExe</OutputType>
-    <TargetFramework>net6.0-windows</TargetFramework>
+    <TargetFramework>net6.0-windows10.0.18362.0</TargetFramework>
     <OutputPath>$(SolutionDir)$(Platform)\$(Configuration)\modules\Awake</OutputPath>
     <Nullable>enable</Nullable>
     <Platforms>x64</Platforms>
@@ -14,6 +14,11 @@
     <AssemblyName>PowerToys.Awake</AssemblyName>
     <Version>$(Version).0</Version>
     <ApplicationIcon>Images\Awake.ico</ApplicationIcon>
+    <SupportedOSPlatformVersion>10.0.18362.0</SupportedOSPlatformVersion>
+    <PackageProjectUrl>https://awake.den.dev</PackageProjectUrl>
+    <RepositoryUrl>https://github.com/microsoft/powertoys</RepositoryUrl>
+    <EnableNETAnalyzers>true</EnableNETAnalyzers>
+    <AnalysisMode>Recommended</AnalysisMode>
   </PropertyGroup>
 
   <PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Debug|x64'">
@@ -34,6 +39,10 @@
     <WarningLevel>4</WarningLevel>
   </PropertyGroup>
 
+  <ItemGroup>
+    <None Remove="Images\Awake.ico" />
+  </ItemGroup>
+
   <ItemGroup>
     <PackageReference Include="Microsoft.CodeAnalysis.NetAnalyzers" Version="6.0.0">
       <PrivateAssets>all</PrivateAssets>
@@ -70,6 +79,11 @@
       <Link>StyleCop.json</Link>
     </AdditionalFiles>
   </ItemGroup>
+  <ItemGroup>
+    <Content Include="Images\Awake.ico">
+      <CopyToOutputDirectory>Always</CopyToOutputDirectory>
+    </Content>
+  </ItemGroup>
   <ItemGroup>
     <PackageReference Include="StyleCop.Analyzers">
       <Version>1.1.118</Version>
@@ -77,9 +91,4 @@
       <PrivateAssets>all</PrivateAssets>
     </PackageReference>
   </ItemGroup>
-  <ItemGroup>
-    <None Update="Images\Awake.ico">
-      <CopyToOutputDirectory>Always</CopyToOutputDirectory>
-    </None>
-  </ItemGroup>
 </Project>

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

@@ -3,8 +3,11 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
 using System.IO;
 using System.Runtime.InteropServices;
+using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 using Awake.Core.Models;
@@ -99,9 +102,16 @@ namespace Awake.Core
             _tokenSource = new CancellationTokenSource();
             _threadToken = _tokenSource.Token;
 
-            _runnerThread = Task.Run(() => RunIndefiniteLoop(keepDisplayOn), _threadToken)
-                .ContinueWith((result) => callback(result.Result), TaskContinuationOptions.OnlyOnRanToCompletion)
-                .ContinueWith((result) => failureCallback, TaskContinuationOptions.NotOnRanToCompletion);
+            try
+            {
+                _runnerThread = Task.Run(() => RunIndefiniteLoop(keepDisplayOn), _threadToken)
+                    .ContinueWith((result) => callback(result.Result), TaskContinuationOptions.OnlyOnRanToCompletion)
+                    .ContinueWith((result) => failureCallback, TaskContinuationOptions.NotOnRanToCompletion);
+            }
+            catch (Exception ex)
+            {
+                _log.Error(ex.Message);
+            }
         }
 
         public static void SetNoKeepAwake()
@@ -181,6 +191,26 @@ namespace Awake.Core
             }
         }
 
+        internal static void CompleteExit(int exitCode, bool force = false)
+        {
+            APIHelper.SetNoKeepAwake();
+            TrayHelper.ClearTray();
+
+            // Because we are running a message loop for the tray, we can't just use Environment.Exit,
+            // but have to make sure that we properly send the termination message.
+            IntPtr windowHandle = APIHelper.GetHiddenWindow();
+
+            if (windowHandle != IntPtr.Zero)
+            {
+                NativeMethods.SendMessage(windowHandle, NativeConstants.WM_CLOSE, 0, string.Empty);
+            }
+
+            if (force)
+            {
+                Environment.Exit(exitCode);
+            }
+        }
+
         private static bool RunTimedLoop(uint seconds, bool keepDisplayOn = true)
         {
             bool success = false;
@@ -262,5 +292,54 @@ namespace Awake.Core
                 return string.Empty;
             }
         }
+
+        [SuppressMessage("Performance", "CA1806:Do not ignore method results", Justification = "Function returns DWORD value that identifies the current thread, but we do not need it.")]
+        public static IEnumerable<IntPtr> EnumerateWindowsForProcess(int processId)
+        {
+            var handles = new List<IntPtr>();
+            IntPtr hCurrentWnd = IntPtr.Zero;
+
+            do
+            {
+                hCurrentWnd = NativeMethods.FindWindowEx(IntPtr.Zero, hCurrentWnd, null, null);
+                NativeMethods.GetWindowThreadProcessId(hCurrentWnd, out uint targetProcessId);
+                if (targetProcessId == processId)
+                {
+                    handles.Add(hCurrentWnd);
+                }
+            }
+            while (hCurrentWnd != IntPtr.Zero);
+
+            return handles;
+        }
+
+        [SuppressMessage("Globalization", "CA1305:Specify IFormatProvider", Justification = "In this context, the string is only converted to a hex value.")]
+        public static IntPtr GetHiddenWindow()
+        {
+            IEnumerable<IntPtr> windowHandles = EnumerateWindowsForProcess(Environment.ProcessId);
+            var domain = AppDomain.CurrentDomain.GetHashCode().ToString("x");
+            string targetClass = $"{InternalConstants.TrayWindowId}{domain}";
+
+            foreach (var handle in windowHandles)
+            {
+                StringBuilder className = new (256);
+                int classQueryResult = NativeMethods.GetClassName(handle, className, className.Capacity);
+                if (classQueryResult != 0 && className.ToString().StartsWith(targetClass, StringComparison.InvariantCultureIgnoreCase))
+                {
+                    return handle;
+                }
+            }
+
+            return IntPtr.Zero;
+        }
+
+        public static Dictionary<string, int> GetDefaultTrayOptions()
+        {
+            Dictionary<string, int> optionsList = new Dictionary<string, int>();
+            optionsList.Add("30 minutes", 1800);
+            optionsList.Add("1 hour", 3600);
+            optionsList.Add("2 hours", 7200);
+            return optionsList;
+        }
     }
 }

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

@@ -0,0 +1,30 @@
+// Copyright (c) Microsoft Corporation
+// The Microsoft Corporation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using System.Collections.Generic;
+
+namespace Awake.Core
+{
+    internal static class ExtensionMethods
+    {
+        public static void AddRange<T>(this ICollection<T> target, IEnumerable<T> source)
+        {
+            if (target == null)
+            {
+                throw new ArgumentNullException(nameof(target));
+            }
+
+            if (source == null)
+            {
+                throw new ArgumentNullException(nameof(source));
+            }
+
+            foreach (var element in source)
+            {
+                target.Add(element);
+            }
+        }
+    }
+}

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

@@ -8,5 +8,6 @@ namespace Awake.Core
     {
         internal const string AppName = "Awake";
         internal const string FullAppName = "PowerToys " + AppName;
+        internal const string TrayWindowId = "WindowsForms10.Window.0.app.0.";
     }
 }

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

@@ -0,0 +1,15 @@
+// Copyright (c) Microsoft Corporation
+// The Microsoft Corporation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+namespace Awake.Core.Models
+{
+    internal enum TrayCommands : uint
+    {
+        TC_DISPLAY_SETTING = NativeConstants.WM_USER + 1,
+        TC_MODE_PASSIVE = NativeConstants.WM_USER + 2,
+        TC_MODE_INDEFINITE = NativeConstants.WM_USER + 3,
+        TC_EXIT = NativeConstants.WM_USER + 4,
+        TC_TIME = NativeConstants.WM_USER + 5,
+    }
+}

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

@@ -0,0 +1,26 @@
+// Copyright (c) Microsoft Corporation
+// The Microsoft Corporation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+#pragma warning disable SA1310 // Field names should not contain underscore
+
+namespace Awake.Core
+{
+    internal class NativeConstants
+    {
+        internal const uint WM_COMMAND = 0x111;
+        internal const uint WM_USER = 0x400;
+        internal const uint WM_GETTEXT = 0x000D;
+        internal const uint WM_CLOSE = 0x0010;
+
+        // Popup menu constants.
+        internal const uint MF_BYPOSITION = 1024;
+        internal const uint MF_STRING = 0;
+        internal const uint MF_MENUBREAK = 0x00000040;
+        internal const uint MF_SEPARATOR = 0x00000800;
+        internal const uint MF_POPUP = 0x00000010;
+        internal const uint MF_UNCHECKED = 0x00000000;
+        internal const uint MF_CHECKED = 0x00000008;
+        internal const uint MF_OWNERDRAW = 0x00000100;
+    }
+}

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

@@ -5,19 +5,22 @@
 using System;
 using System.IO;
 using System.Runtime.InteropServices;
+using System.Text;
 using Awake.Core.Models;
 
 namespace Awake.Core
 {
     internal static class NativeMethods
     {
+        internal delegate bool EnumThreadDelegate(IntPtr hWnd, IntPtr lParam);
+
         [DllImport("Powrprof.dll", SetLastError = true)]
         internal static extern bool GetPwrCapabilities(out SystemPowerCapabilities lpSystemPowerCapabilities);
 
         [DllImport("kernel32.dll", SetLastError = true)]
         internal static extern bool SetConsoleCtrlHandler(ConsoleEventHandler handler, bool add);
 
-        [DllImport("kernel32.dll", CharSet = CharSet.Auto, SetLastError = true)]
+        [DllImport("kernel32.dll", SetLastError = true, CharSet = CharSet.Auto)]
         internal static extern ExecutionState SetThreadExecutionState(ExecutionState esFlags);
 
         [DllImport("kernel32.dll", SetLastError = true)]
@@ -27,10 +30,10 @@ namespace Awake.Core
         [DllImport("kernel32.dll", SetLastError = true)]
         internal static extern bool SetStdHandle(int nStdHandle, IntPtr hHandle);
 
-        [DllImport("kernel32.dll")]
+        [DllImport("kernel32.dll", SetLastError = true)]
         internal static extern uint GetCurrentThreadId();
 
-        [DllImport("kernel32.dll", CharSet = CharSet.Auto, SetLastError = true)]
+        [DllImport("kernel32.dll", SetLastError = true, CharSet = CharSet.Auto)]
         internal static extern IntPtr CreateFile(
         [MarshalAs(UnmanagedType.LPWStr)] string filename,
         [MarshalAs(UnmanagedType.U4)] uint access,
@@ -39,5 +42,33 @@ namespace Awake.Core
         [MarshalAs(UnmanagedType.U4)] FileMode creationDisposition,
         [MarshalAs(UnmanagedType.U4)] FileAttributes flagsAndAttributes,
         IntPtr templateFile);
+
+        [DllImport("user32.dll", SetLastError = true)]
+        internal static extern IntPtr CreatePopupMenu();
+
+        [DllImport("user32.dll", SetLastError = true, CharSet = CharSet.Unicode)]
+        internal static extern bool InsertMenu(IntPtr hMenu, uint uPosition, uint uFlags, uint uIDNewItem, string lpNewItem);
+
+        [DllImport("user32.dll", SetLastError = true)]
+        internal static extern bool TrackPopupMenuEx(IntPtr hmenu, uint fuFlags, int x, int y, IntPtr hwnd, IntPtr lptpm);
+
+        [DllImport("user32.dll", SetLastError = true, CharSet = CharSet.Unicode)]
+        internal static extern int GetClassName(IntPtr hWnd, StringBuilder lpClassName, int nMaxCount);
+
+        [DllImport("user32.dll", SetLastError = true, CharSet = CharSet.Unicode)]
+        internal static extern IntPtr FindWindowEx(IntPtr parentHandle, IntPtr hWndChildAfter, string? className, string? windowTitle);
+
+        [DllImport("user32.dll", SetLastError = true)]
+        internal static extern uint GetWindowThreadProcessId(IntPtr hWnd, out uint lpdwProcessId);
+
+        [DllImport("user32.dll", SetLastError = true)]
+        [return: MarshalAs(UnmanagedType.Bool)]
+        internal static extern bool SetForegroundWindow(IntPtr hWnd);
+
+        [DllImport("user32.dll", SetLastError = true, CharSet = CharSet.Unicode)]
+        internal static extern IntPtr SendMessage(IntPtr hWnd, uint Msg, nuint wParam, string lParam);
+
+        [DllImport("user32.dll", SetLastError = true)]
+        internal static extern bool DestroyMenu(IntPtr hMenu);
     }
 }

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

@@ -3,11 +3,13 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
 using System.Drawing;
-using System.IO;
-using System.Text.Json;
+using System.Linq;
 using System.Threading.Tasks;
 using System.Windows.Forms;
+using Awake.Core.Models;
 using Microsoft.PowerToys.Settings.UI.Library;
 using NLog;
 
@@ -20,19 +22,18 @@ namespace Awake.Core
     {
         private static readonly Logger _log;
 
+        private static IntPtr _trayMenu;
+
+        private static IntPtr TrayMenu { get => _trayMenu; set => _trayMenu = value; }
+
         private static NotifyIcon? _trayIcon;
 
         private static NotifyIcon TrayIcon { get => _trayIcon; set => _trayIcon = value; }
 
-        private static SettingsUtils? _moduleSettings;
-
-        private static SettingsUtils ModuleSettings { get => _moduleSettings; set => _moduleSettings = value; }
-
         static TrayHelper()
         {
             _log = LogManager.GetCurrentClassLogger();
             TrayIcon = new NotifyIcon();
-            ModuleSettings = new SettingsUtils();
         }
 
         public static void InitializeTray(string text, Icon icon, ContextMenuStrip? contextMenu = null)
@@ -40,17 +41,48 @@ namespace Awake.Core
             Task.Factory.StartNew(
                 (tray) =>
                 {
-                    ((NotifyIcon?)tray).Text = text;
-                    ((NotifyIcon?)tray).Icon = icon;
-                    ((NotifyIcon?)tray).ContextMenuStrip = contextMenu;
-                    ((NotifyIcon?)tray).Visible = true;
-
-                    _log.Info("Setting up the tray.");
-                    Application.Run();
-                    _log.Info("Tray setup complete.");
+                    try
+                    {
+                        _log.Info("Setting up the tray.");
+                        ((NotifyIcon?)tray).Text = text;
+                        ((NotifyIcon?)tray).Icon = icon;
+                        ((NotifyIcon?)tray).ContextMenuStrip = contextMenu;
+                        ((NotifyIcon?)tray).Visible = true;
+                        ((NotifyIcon?)tray).MouseClick += TrayClickHandler;
+                        Application.AddMessageFilter(new TrayMessageFilter());
+                        Application.Run();
+                        _log.Info("Tray setup complete.");
+                    }
+                    catch (Exception ex)
+                    {
+                        _log.Error($"An error occurred initializing the tray. {ex.Message}");
+                        _log.Error($"{ex.StackTrace}");
+                    }
                 }, TrayIcon);
         }
 
+        /// <summary>
+        /// Function used to construct the context menu in the tray natively.
+        /// </summary>
+        /// <remarks>
+        /// We need to use the Windows API here instead of the common control exposed
+        /// by NotifyIcon because the one that is built into the Windows Forms stack
+        /// hasn't been updated in a while and is looking like Office XP. That introduces
+        /// scalability and coloring changes on any OS past Windows XP.
+        /// </remarks>
+        /// <param name="sender">The sender that triggers the handler.</param>
+        /// <param name="e">MouseEventArgs instance containing mouse click event information.</param>
+        private static void TrayClickHandler(object? sender, MouseEventArgs e)
+        {
+            IntPtr windowHandle = APIHelper.GetHiddenWindow();
+
+            if (windowHandle != IntPtr.Zero)
+            {
+                NativeMethods.SetForegroundWindow(windowHandle);
+                NativeMethods.TrackPopupMenuEx(TrayMenu, 0, Cursor.Position.X, Cursor.Position.Y, windowHandle, IntPtr.Zero);
+            }
+        }
+
         public static void ClearTray()
         {
             TrayIcon.Icon = null;
@@ -63,227 +95,48 @@ namespace Awake.Core
                 text,
                 settings.Properties.KeepDisplayOn,
                 settings.Properties.Mode,
-                PassiveKeepAwakeCallback(InternalConstants.AppName),
-                IndefiniteKeepAwakeCallback(InternalConstants.AppName),
-                TimedKeepAwakeCallback(InternalConstants.AppName),
-                KeepDisplayOnCallback(InternalConstants.AppName),
-                ExitCallback());
+                settings.Properties.TrayTimeShortcuts);
         }
 
-        private static Action ExitCallback()
+        [SuppressMessage("StyleCop.CSharp.SpacingRules", "SA1005:Single line comments should begin with single space", Justification = "For debugging purposes - will remove later.")]
+        public static void SetTray(string text, bool keepDisplayOn, AwakeMode mode, Dictionary<string, int> trayTimeShortcuts)
         {
-            return () =>
+            if (TrayMenu != IntPtr.Zero)
             {
-                Environment.Exit(Environment.ExitCode);
-            };
-        }
-
-        private static Action KeepDisplayOnCallback(string moduleName)
-        {
-            return () =>
-            {
-                AwakeSettings currentSettings;
-
-                try
+                var destructionStatus = NativeMethods.DestroyMenu(TrayMenu);
+                if (destructionStatus != true)
                 {
-                    currentSettings = ModuleSettings.GetSettings<AwakeSettings>(moduleName);
-                }
-                catch (FileNotFoundException)
-                {
-                    currentSettings = new AwakeSettings();
+                    _log.Error("Failed to destroy tray menu and free up memory.");
                 }
+            }
 
-                currentSettings.Properties.KeepDisplayOn = !currentSettings.Properties.KeepDisplayOn;
+            TrayMenu = NativeMethods.CreatePopupMenu();
+            NativeMethods.InsertMenu(TrayMenu, 0, NativeConstants.MF_BYPOSITION | NativeConstants.MF_STRING, (uint)TrayCommands.TC_EXIT, "Exit");
+            NativeMethods.InsertMenu(TrayMenu, 0, NativeConstants.MF_BYPOSITION | NativeConstants.MF_SEPARATOR, 0, string.Empty);
+            NativeMethods.InsertMenu(TrayMenu, 0, NativeConstants.MF_BYPOSITION | NativeConstants.MF_STRING | (keepDisplayOn ? NativeConstants.MF_CHECKED : NativeConstants.MF_UNCHECKED), (uint)TrayCommands.TC_DISPLAY_SETTING, "Keep screen on");
 
-                ModuleSettings.SaveSettings(JsonSerializer.Serialize(currentSettings), moduleName);
-            };
-        }
-
-        private static Action<uint, uint> TimedKeepAwakeCallback(string moduleName)
-        {
-            return (hours, minutes) =>
+            // In case there are no tray shortcuts defined for the application default to a
+            // reasonable initial set.
+            if (trayTimeShortcuts.Count == 0)
             {
-                AwakeSettings currentSettings;
+                trayTimeShortcuts.AddRange(APIHelper.GetDefaultTrayOptions());
+            }
 
-                try
-                {
-                    currentSettings = ModuleSettings.GetSettings<AwakeSettings>(moduleName);
-                }
-                catch (FileNotFoundException)
-                {
-                    currentSettings = new AwakeSettings();
-                }
-
-                currentSettings.Properties.Mode = AwakeMode.TIMED;
-                currentSettings.Properties.Hours = hours;
-                currentSettings.Properties.Minutes = minutes;
-
-                ModuleSettings.SaveSettings(JsonSerializer.Serialize(currentSettings), moduleName);
-            };
-        }
-
-        private static Action PassiveKeepAwakeCallback(string moduleName)
-        {
-            return () =>
+            // TODO: Make sure that this loads from JSON instead of being hard-coded.
+            var awakeTimeMenu = NativeMethods.CreatePopupMenu();
+            for (int i = 0; i < trayTimeShortcuts.Count; i++)
             {
-                AwakeSettings currentSettings;
+                NativeMethods.InsertMenu(awakeTimeMenu, (uint)i, NativeConstants.MF_BYPOSITION | NativeConstants.MF_STRING, (uint)TrayCommands.TC_TIME + (uint)i, trayTimeShortcuts.ElementAt(i).Key);
+            }
 
-                try
-                {
-                    currentSettings = ModuleSettings.GetSettings<AwakeSettings>(moduleName);
-                }
-                catch (FileNotFoundException)
-                {
-                    currentSettings = new AwakeSettings();
-                }
+            var modeMenu = NativeMethods.CreatePopupMenu();
+            NativeMethods.InsertMenu(modeMenu, 0, NativeConstants.MF_BYPOSITION | NativeConstants.MF_STRING | (mode == AwakeMode.PASSIVE ? NativeConstants.MF_CHECKED : NativeConstants.MF_UNCHECKED), (uint)TrayCommands.TC_MODE_PASSIVE, "Off (keep using the selected power plan)");
+            NativeMethods.InsertMenu(modeMenu, 1, NativeConstants.MF_BYPOSITION | NativeConstants.MF_STRING | (mode == AwakeMode.INDEFINITE ? NativeConstants.MF_CHECKED : NativeConstants.MF_UNCHECKED), (uint)TrayCommands.TC_MODE_INDEFINITE, "Keep awake indefinitely");
 
-                currentSettings.Properties.Mode = AwakeMode.PASSIVE;
-
-                ModuleSettings.SaveSettings(JsonSerializer.Serialize(currentSettings), moduleName);
-            };
-        }
-
-        private static Action IndefiniteKeepAwakeCallback(string moduleName)
-        {
-            return () =>
-            {
-                AwakeSettings currentSettings;
-
-                try
-                {
-                    currentSettings = ModuleSettings.GetSettings<AwakeSettings>(moduleName);
-                }
-                catch (FileNotFoundException)
-                {
-                    currentSettings = new AwakeSettings();
-                }
-
-                currentSettings.Properties.Mode = AwakeMode.INDEFINITE;
-
-                ModuleSettings.SaveSettings(JsonSerializer.Serialize(currentSettings), moduleName);
-            };
-        }
-
-        public static void SetTray(string text, bool keepDisplayOn, AwakeMode mode, Action passiveKeepAwakeCallback, Action indefiniteKeepAwakeCallback, Action<uint, uint> timedKeepAwakeCallback, Action keepDisplayOnCallback, Action exitCallback)
-        {
-            ContextMenuStrip? contextMenuStrip = new ContextMenuStrip();
-
-            // Main toolstrip.
-            ToolStripMenuItem? operationContextMenu = new ToolStripMenuItem
-            {
-                Text = "Mode",
-            };
-
-            // No keep-awake menu item.
-            CheckButtonToolStripMenuItem? passiveMenuItem = new CheckButtonToolStripMenuItem
-            {
-                Text = "Off (Keep using the selected power plan)",
-            };
-
-            passiveMenuItem.Checked = mode == AwakeMode.PASSIVE;
-
-            passiveMenuItem.Click += (e, s) =>
-            {
-                // User opted to set the mode to indefinite, so we need to write new settings.
-                passiveKeepAwakeCallback();
-            };
-
-            // Indefinite keep-awake menu item.
-            CheckButtonToolStripMenuItem? indefiniteMenuItem = new CheckButtonToolStripMenuItem
-            {
-                Text = "Keep awake indefinitely",
-            };
-
-            indefiniteMenuItem.Checked = mode == AwakeMode.INDEFINITE;
-
-            indefiniteMenuItem.Click += (e, s) =>
-            {
-                // User opted to set the mode to indefinite, so we need to write new settings.
-                indefiniteKeepAwakeCallback();
-            };
-
-            CheckButtonToolStripMenuItem? displayOnMenuItem = new CheckButtonToolStripMenuItem
-            {
-                Text = "Keep screen on",
-            };
-
-            displayOnMenuItem.Checked = keepDisplayOn;
-
-            displayOnMenuItem.Click += (e, s) =>
-            {
-                // User opted to set the display mode directly.
-                keepDisplayOnCallback();
-            };
-
-            // Timed keep-awake menu item
-            ToolStripMenuItem? timedMenuItem = new ToolStripMenuItem
-            {
-                Text = "Keep awake temporarily",
-            };
-
-            timedMenuItem.Checked = mode == AwakeMode.TIMED;
-            timedMenuItem.AccessibleName = timedMenuItem.Text + (timedMenuItem.Checked ? ". Checked. " : ". UnChecked. ");
-
-            ToolStripMenuItem? halfHourMenuItem = new ToolStripMenuItem
-            {
-                Text = "30 minutes",
-            };
-
-            halfHourMenuItem.Click += (e, s) =>
-            {
-                // User is setting the keep-awake to 30 minutes.
-                timedKeepAwakeCallback(0, 30);
-            };
-
-            ToolStripMenuItem? oneHourMenuItem = new ToolStripMenuItem
-            {
-                Text = "1 hour",
-            };
-
-            oneHourMenuItem.Click += (e, s) =>
-            {
-                // User is setting the keep-awake to 1 hour.
-                timedKeepAwakeCallback(1, 0);
-            };
-
-            ToolStripMenuItem? twoHoursMenuItem = new ToolStripMenuItem
-            {
-                Text = "2 hours",
-            };
-
-            twoHoursMenuItem.Click += (e, s) =>
-            {
-                // User is setting the keep-awake to 2 hours.
-                timedKeepAwakeCallback(2, 0);
-            };
-
-            // Exit menu item.
-            ToolStripMenuItem? exitContextMenu = new ToolStripMenuItem
-            {
-                Text = "Exit",
-            };
-
-            exitContextMenu.Click += (e, s) =>
-            {
-                // User is setting the keep-awake to 2 hours.
-                exitCallback();
-            };
-
-            timedMenuItem.DropDownItems.Add(halfHourMenuItem);
-            timedMenuItem.DropDownItems.Add(oneHourMenuItem);
-            timedMenuItem.DropDownItems.Add(twoHoursMenuItem);
-
-            operationContextMenu.DropDownItems.Add(passiveMenuItem);
-            operationContextMenu.DropDownItems.Add(indefiniteMenuItem);
-            operationContextMenu.DropDownItems.Add(timedMenuItem);
-
-            contextMenuStrip.Items.Add(operationContextMenu);
-            contextMenuStrip.Items.Add(displayOnMenuItem);
-            contextMenuStrip.Items.Add(new ToolStripSeparator());
-            contextMenuStrip.Items.Add(exitContextMenu);
+            NativeMethods.InsertMenu(modeMenu, 2, NativeConstants.MF_BYPOSITION | NativeConstants.MF_POPUP | (mode == AwakeMode.TIMED ? NativeConstants.MF_CHECKED : NativeConstants.MF_UNCHECKED), (uint)awakeTimeMenu, "Keep awake temporarily");
+            NativeMethods.InsertMenu(TrayMenu, 0, NativeConstants.MF_BYPOSITION | NativeConstants.MF_POPUP, (uint)modeMenu, "Mode");
 
             TrayIcon.Text = text;
-            TrayIcon.ContextMenuStrip = contextMenuStrip;
         }
 
         private class CheckButtonToolStripMenuItemAccessibleObject : ToolStripItem.ToolStripItemAccessibleObject

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

@@ -0,0 +1,152 @@
+// Copyright (c) Microsoft Corporation
+// The Microsoft Corporation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using System.IO;
+using System.Linq;
+using System.Text.Json;
+using System.Windows.Forms;
+using Awake.Core.Models;
+using Microsoft.PowerToys.Settings.UI.Library;
+
+#pragma warning disable CS8603 // Possible null reference return.
+
+namespace Awake.Core
+{
+    public class TrayMessageFilter : IMessageFilter
+    {
+        private static SettingsUtils? _moduleSettings;
+
+        private static SettingsUtils ModuleSettings { get => _moduleSettings; set => _moduleSettings = value; }
+
+        public TrayMessageFilter()
+        {
+            ModuleSettings = new SettingsUtils();
+        }
+
+        public bool PreFilterMessage(ref Message m)
+        {
+            var trayCommandsSize = Enum.GetNames(typeof(TrayCommands)).Length;
+
+            switch (m.Msg)
+            {
+                case (int)NativeConstants.WM_COMMAND:
+                    var targetCommandIndex = m.WParam.ToInt64() & 0xFFFF;
+                    switch (targetCommandIndex)
+                    {
+                        case (long)TrayCommands.TC_EXIT:
+                            ExitCommandHandler();
+                            break;
+                        case (long)TrayCommands.TC_DISPLAY_SETTING:
+                            DisplaySettingCommandHandler(InternalConstants.AppName);
+                            break;
+                        case (long)TrayCommands.TC_MODE_INDEFINITE:
+                            IndefiniteKeepAwakeCommandHandler(InternalConstants.AppName);
+                            break;
+                        case (long)TrayCommands.TC_MODE_PASSIVE:
+                            PassiveKeepAwakeCommandHandler(InternalConstants.AppName);
+                            break;
+                        case var _ when targetCommandIndex >= trayCommandsSize:
+                            // Format for the timer block:
+                            // TrayCommands.TC_TIME + ZERO_BASED_INDEX_IN_SETTINGS
+                            AwakeSettings settings = ModuleSettings.GetSettings<AwakeSettings>(InternalConstants.AppName);
+                            if (settings.Properties.TrayTimeShortcuts.Count == 0)
+                            {
+                                settings.Properties.TrayTimeShortcuts.AddRange(APIHelper.GetDefaultTrayOptions());
+                            }
+
+                            int index = (int)targetCommandIndex - (int)TrayCommands.TC_TIME;
+                            var targetTime = settings.Properties.TrayTimeShortcuts.ElementAt(index).Value;
+                            TimedKeepAwakeCommandHandler(InternalConstants.AppName, targetTime);
+                            break;
+                    }
+
+                    break;
+            }
+
+            return false;
+        }
+
+        private static void ExitCommandHandler()
+        {
+            APIHelper.CompleteExit(0, true);
+        }
+
+        private static void DisplaySettingCommandHandler(string moduleName)
+        {
+            AwakeSettings currentSettings;
+
+            try
+            {
+                currentSettings = ModuleSettings.GetSettings<AwakeSettings>(moduleName);
+            }
+            catch (FileNotFoundException)
+            {
+                currentSettings = new AwakeSettings();
+            }
+
+            currentSettings.Properties.KeepDisplayOn = !currentSettings.Properties.KeepDisplayOn;
+
+            ModuleSettings.SaveSettings(JsonSerializer.Serialize(currentSettings), moduleName);
+        }
+
+        private static void TimedKeepAwakeCommandHandler(string moduleName, int seconds)
+        {
+            TimeSpan timeSpan = TimeSpan.FromSeconds(seconds);
+
+            AwakeSettings currentSettings;
+
+            try
+            {
+                currentSettings = ModuleSettings.GetSettings<AwakeSettings>(moduleName);
+            }
+            catch (FileNotFoundException)
+            {
+                currentSettings = new AwakeSettings();
+            }
+
+            currentSettings.Properties.Mode = AwakeMode.TIMED;
+            currentSettings.Properties.Hours = (uint)timeSpan.Hours;
+            currentSettings.Properties.Minutes = (uint)timeSpan.Minutes;
+
+            ModuleSettings.SaveSettings(JsonSerializer.Serialize(currentSettings), moduleName);
+        }
+
+        private static void PassiveKeepAwakeCommandHandler(string moduleName)
+        {
+            AwakeSettings currentSettings;
+
+            try
+            {
+                currentSettings = ModuleSettings.GetSettings<AwakeSettings>(moduleName);
+            }
+            catch (FileNotFoundException)
+            {
+                currentSettings = new AwakeSettings();
+            }
+
+            currentSettings.Properties.Mode = AwakeMode.PASSIVE;
+
+            ModuleSettings.SaveSettings(JsonSerializer.Serialize(currentSettings), moduleName);
+        }
+
+        private static void IndefiniteKeepAwakeCommandHandler(string moduleName)
+        {
+            AwakeSettings currentSettings;
+
+            try
+            {
+                currentSettings = ModuleSettings.GetSettings<AwakeSettings>(moduleName);
+            }
+            catch (FileNotFoundException)
+            {
+                currentSettings = new AwakeSettings();
+            }
+
+            currentSettings.Properties.Mode = AwakeMode.INDEFINITE;
+
+            ModuleSettings.SaveSettings(JsonSerializer.Serialize(currentSettings), moduleName);
+        }
+    }
+}

src/modules/awake/Awake/NLog.config

@@ -2,7 +2,7 @@
 <nlog xmlns="http://www.nlog-project.org/schemas/NLog.xsd"
       xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance">
 
-  <variable name="buildId" value="ARBITER_01312022" />
+  <variable name="buildId" value="LIBRARIAN_03202022" />
   
   <targets async="true">
     <target name="logfile"

src/modules/awake/Awake/Program.cs

@@ -3,6 +3,7 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Collections.Generic;
 using System.CommandLine;
 using System.CommandLine.Invocation;
 using System.Diagnostics;
@@ -14,6 +15,7 @@ using System.Reactive.Linq;
 using System.Reflection;
 using System.Text.Json;
 using System.Threading;
+using System.Threading.Tasks;
 using Awake.Core;
 using Awake.Core.Models;
 using interop;
@@ -72,6 +74,12 @@ namespace Awake
             _log.Info($"OS: {Environment.OSVersion}");
             _log.Info($"OS Build: {APIHelper.GetOperatingSystemBuild()}");
 
+            TaskScheduler.UnobservedTaskException += (sender, args) =>
+            {
+                Trace.WriteLine($"Task scheduler error: {args.Exception.Message}"); // somebody forgot to check!
+                args.SetObserved();
+            };
+
             // To make it easier to diagnose future issues, let's get the
             // system power capabilities and aggregate them in the log.
             NativeMethods.GetPwrCapabilities(out _powerCapabilities);
@@ -159,18 +167,7 @@ namespace Awake
         {
             _log.Info(message);
 
-            APIHelper.SetNoKeepAwake();
-            TrayHelper.ClearTray();
-
-            // Because we are running a message loop for the tray, we can't just use Environment.Exit,
-            // but have to make sure that we properly send the termination message.
-            bool cwResult = System.Diagnostics.Process.GetCurrentProcess().CloseMainWindow();
-            _log.Info($"Request to close main window status: {cwResult}");
-
-            if (force)
-            {
-                Environment.Exit(exitCode);
-            }
+            APIHelper.CompleteExit(exitCode, force);
         }
 
         private static void HandleCommandLineArguments(bool usePtConfig, bool displayOn, uint timeLimit, int pid)
@@ -204,7 +201,7 @@ namespace Awake
                         }
                     }).Start();
 
-                    TrayHelper.InitializeTray(InternalConstants.FullAppName, new Icon("modules/Awake/Images/Awake.ico"));
+                    TrayHelper.InitializeTray(InternalConstants.FullAppName, new Icon("modules/awake/images/awake.ico"));
 
                     string? settingsPath = _settingsUtils.GetSettingsFilePath(InternalConstants.AppName);
                     _log.Info($"Reading configuration file: {settingsPath}");
@@ -293,6 +290,8 @@ namespace Awake
 
                 if (settings != null)
                 {
+                    _log.Info($"Identified custom time shortcuts for the tray: {settings.Properties.TrayTimeShortcuts.Count}");
+
                     switch (settings.Properties.Mode)
                     {
                         case AwakeMode.PASSIVE:

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

@@ -2,6 +2,7 @@
 // The Microsoft Corporation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+using System.Collections.Generic;
 using System.Text.Json.Serialization;
 
 namespace Microsoft.PowerToys.Settings.UI.Library
@@ -14,6 +15,7 @@ namespace Microsoft.PowerToys.Settings.UI.Library
             Mode = AwakeMode.PASSIVE;
             Hours = 0;
             Minutes = 0;
+            TrayTimeShortcuts = new Dictionary<string, int>();
         }
 
         [JsonPropertyName("awake_keep_display_on")]
@@ -27,6 +29,9 @@ namespace Microsoft.PowerToys.Settings.UI.Library
 
         [JsonPropertyName("awake_minutes")]
         public uint Minutes { get; set; }
+
+        [JsonPropertyName("tray_times")]
+        public Dictionary<string, int> TrayTimeShortcuts { get; set; }
     }
 
     public enum AwakeMode