Register restart for resume

schemas/JSON/settings/settings.schema.0.2.json

@@ -148,6 +148,12 @@
           "description": "Default install location to use for packages that require it when not specified",
           "type": "string",
           "maxLength": "32767"
+        },
+        "maxResumes": {
+          "description": "The maximum number of resumes allowed for a single resume id. This is to prevent continuous reboots if an install that requires a reboot is not properly detected.",
+          "type": "integer",
+          "default": 3,
+          "minimum": 1
         }
       }
     },

src/AppInstallerCLICore/Argument.cpp

@@ -185,6 +185,8 @@ namespace AppInstaller::CLI
         // Resume command
         case Execution::Args::Type::ResumeId:
             return { type, "resume-id"_liv, 'g', ArgTypeCategory::None };
+        case Execution::Args::Type::IgnoreResumeLimit:
+            return { type, "ignore-resume-limit"_liv, ArgTypeCategory::None };
 
         // Configuration commands
         case Execution::Args::Type::ConfigurationFile:
@@ -365,6 +367,8 @@ namespace AppInstaller::CLI
             return Argument{ type, Resource::String::ResumeIdArgumentDescription, ArgumentType::Standard, true };
         case Args::Type::AllowReboot:
             return Argument{ type, Resource::String::AllowRebootArgumentDescription, ArgumentType::Flag, ExperimentalFeature::Feature::Reboot };
+        case Args::Type::IgnoreResumeLimit:
+            return Argument{ type, Resource::String::IgnoreResumeLimitArgumentDescription, ArgumentType::Flag, ExperimentalFeature::Feature::Resume };
         default:
             THROW_HR(E_UNEXPECTED);
         }

src/AppInstallerCLICore/CheckpointManager.cpp

@@ -82,6 +82,8 @@ namespace AppInstaller::Checkpoints
                 automaticCheckpoint.SetMany(AutomaticCheckpointData::Arguments, argument, values);
             }
         }
+
+        automaticCheckpoint.Set(AutomaticCheckpointData::ResumeCount, {}, std::to_string(0));
     }
 
     void LoadCommandArgsFromAutomaticCheckpoint(CLI::Execution::Context& context, Checkpoint<AutomaticCheckpointData>& automaticCheckpoint)

src/AppInstallerCLICore/CheckpointManager.h

@@ -30,6 +30,9 @@ namespace AppInstaller::Checkpoints
         // Gets the file path of the checkpoint database.
         static std::filesystem::path GetCheckpointDatabasePath(const std::string_view& resumeId, bool createCheckpointDirectory = false);
 
+        // Gets the resume id.
+        std::string GetResumeId() { return m_resumeId; };
+
         // Gets the automatic checkpoint.
         std::optional<Checkpoint<AutomaticCheckpointData>> GetAutomaticCheckpoint();
 

src/AppInstallerCLICore/Command.cpp

@@ -6,6 +6,7 @@
 #include <winget/UserSettings.h>
 #include <AppInstallerRuntime.h>
 #include <winget/Locale.h>
+#include <winget/Reboot.h>
 
 using namespace std::string_view_literals;
 using namespace AppInstaller::Utility::literals;
@@ -870,16 +871,47 @@ namespace AppInstaller::CLI
             ExecuteInternal(context);
         }
 
-        if (context.Args.Contains(Execution::Args::Type::OpenLogs))
-        {   
-            // TODO: Consider possibly adding functionality that if the context contains 'Execution::Args::Type::Log' to open the path provided for the log
-            // The above was omitted initially as a security precaution to ensure that user input to '--log' wouldn't be passed directly to ShellExecute
-            ShellExecute(NULL, NULL, Runtime::GetPathTo(Runtime::PathName::DefaultLogLocation).wstring().c_str(), NULL, NULL, SW_SHOWNORMAL);
-        }
+        // NOTE: Reboot logic will still run even if the context is terminated (not including unhandled exceptions).
+        if (Settings::ExperimentalFeature::IsEnabled(Settings::ExperimentalFeature::Feature::Reboot) &&
+            context.Args.Contains(Execution::Args::Type::AllowReboot) &&
+            WI_IsFlagSet(context.GetFlags(), Execution::ContextFlag::RebootRequired))
+        {
+            context.Reporter.Warn() << Resource::String::InitiatingReboot << std::endl;
+
+            if (context.Args.Contains(Execution::Args::Type::Wait))
+            {
+                context.Reporter.PromptForEnter();
+            }
+
+            if (WI_IsFlagSet(context.GetFlags(), Execution::ContextFlag::RegisterResume))
+            {
+                // RegisterResume context flag assumes we already wrote to the RunOnce registry.
+                // Since we are about to initiate a restart, this is no longer needed as a safety net.
+                Reboot::UnregisterRestartForWER();
+
+                context.ClearFlags(Execution::ContextFlag::RegisterResume);
+            }
+
+            context.ClearFlags(Execution::ContextFlag::RebootRequired);
 
-        if (context.Args.Contains(Execution::Args::Type::Wait))
+            if (!Reboot::InitiateReboot())
+            {
+                context.Reporter.Error() << Resource::String::FailedToInitiateReboot << std::endl;
+            }
+        }
+        else
         {
-            context.Reporter.PromptForEnter();
+            if (context.Args.Contains(Execution::Args::Type::OpenLogs))
+            {
+                // TODO: Consider possibly adding functionality that if the context contains 'Execution::Args::Type::Log' to open the path provided for the log
+                // The above was omitted initially as a security precaution to ensure that user input to '--log' wouldn't be passed directly to ShellExecute
+                ShellExecute(NULL, NULL, Runtime::GetPathTo(Runtime::PathName::DefaultLogLocation).wstring().c_str(), NULL, NULL, SW_SHOWNORMAL);
+            }
+
+            if (context.Args.Contains(Execution::Args::Type::Wait))
+            {
+                context.Reporter.PromptForEnter();
+            }
         }
     }
 

src/AppInstallerCLICore/Commands/ResumeCommand.cpp

@@ -54,6 +54,7 @@ namespace AppInstaller::CLI
     {
         return {
             Argument::ForType(Execution::Args::Type::ResumeId),
+            Argument::ForType(Execution::Args::Type::IgnoreResumeLimit),
         };
     }
 
@@ -99,6 +100,19 @@ namespace AppInstaller::CLI
             AICLI_TERMINATE_CONTEXT(APPINSTALLER_CLI_ERROR_CLIENT_VERSION_MISMATCH);
         }
 
+        const auto& resumeCountString = automaticCheckpoint.Get(AutomaticCheckpointData::ResumeCount, {});
+        int resumeCount = std::stoi(resumeCountString);
+        if (!context.Args.Contains(Execution::Args::Type::IgnoreResumeLimit) && resumeCount >= Settings::User().Get<Settings::Setting::MaxResumes>())
+        {
+            std::string manualResumeString = "winget resume -g " + std::string{ resumeId } + " --ignore-resume-limit";
+            context.Reporter.Error() << Resource::String::ResumeLimitExceeded(Utility::LocIndView{ resumeCountString }) << Utility::LocIndView{ manualResumeString } << std::endl;
+            AICLI_TERMINATE_CONTEXT(APPINSTALLER_CLI_ERROR_RESUME_LIMIT_EXCEEDED);
+        }
+        else
+        {
+            automaticCheckpoint.Update(AutomaticCheckpointData::ResumeCount, {}, std::to_string(resumeCount + 1));
+        }
+
         const auto& checkpointCommand = automaticCheckpoint.Get(AutomaticCheckpointData::Command, {});
         AICLI_LOG(CLI, Info, << "Resuming command: " << checkpointCommand);
         std::unique_ptr<Command> commandToResume = FindCommandToResume(checkpointCommand);

src/AppInstallerCLICore/ExecutionArgs.h

@@ -112,6 +112,7 @@ namespace AppInstaller::CLI::Execution
 
             // Resume Command
             ResumeId,
+            IgnoreResumeLimit,
 
             // Configuration
             ConfigurationFile,

src/AppInstallerCLICore/ExecutionContext.cpp

@@ -1,13 +1,14 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 #include "pch.h"
-#include "ExecutionContext.h"
-#include "COMContext.h"
-#include "Argument.h"
-#include "winget/UserSettings.h"
 #include "AppInstallerRuntime.h"
+#include "Argument.h"
+#include "COMContext.h"
 #include "Command.h"
+#include "ExecutionContext.h"
 #include "Public/winget/Checkpoint.h"
+#include "winget/Reboot.h"
+#include "winget/UserSettings.h"
 
 using namespace AppInstaller::Checkpoints;
 
@@ -272,15 +273,29 @@ namespace AppInstaller::CLI::Execution
                 SignalTerminationHandler::Instance().RemoveContext(context);
             }
         }
+
+        bool ShouldRemoveCheckpointDatabase(HRESULT hr)
+        {
+            switch (hr)
+            {
+            case APPINSTALLER_CLI_ERROR_INSTALL_REBOOT_REQUIRED_FOR_INSTALL:
+            case APPINSTALLER_CLI_ERROR_RESUME_LIMIT_EXCEEDED:
+            case APPINSTALLER_CLI_ERROR_CLIENT_VERSION_MISMATCH:
+                return false;
+            default:
+                return true;
+            }
+        }
     }
 
     Context::~Context()
     {
-        if (Settings::ExperimentalFeature::IsEnabled(ExperimentalFeature::Feature::Resume) && !IsTerminated())
+        if (Settings::ExperimentalFeature::IsEnabled(ExperimentalFeature::Feature::Resume))
         {
-            if (m_checkpointManager)
+            if (m_checkpointManager && (!IsTerminated() || ShouldRemoveCheckpointDatabase(GetTerminationHR())))
             {
                 m_checkpointManager->CleanUpDatabase();
+                AppInstaller::Reboot::UnregisterRestartForWER();
             }
         }
 
@@ -448,6 +463,11 @@ namespace AppInstaller::CLI::Execution
     }
 #endif
 
+    std::string Context::GetResumeId()
+    {
+        return m_checkpointManager->GetResumeId();
+    }
+
     std::optional<Checkpoint<AutomaticCheckpointData>> Context::LoadCheckpoint(const std::string& resumeId)
     {
         m_checkpointManager = std::make_unique<AppInstaller::Checkpoints::CheckpointManager>(resumeId);
@@ -468,6 +488,9 @@ namespace AppInstaller::CLI::Execution
         {
             m_checkpointManager = std::make_unique<AppInstaller::Checkpoints::CheckpointManager>();
             m_checkpointManager->CreateAutomaticCheckpoint(*this);
+
+            // Register for restart only when we first call checkpoint to support restarting from an unexpected shutdown.
+            AppInstaller::Reboot::RegisterRestartForWER("resume -g " + GetResumeId());
         }
 
         // TODO: Capture context data for checkpoint.

src/AppInstallerCLICore/ExecutionContext.h

@@ -172,6 +172,9 @@ namespace AppInstaller::CLI::Execution
         bool ShouldExecuteWorkflowTask(const Workflow::WorkflowTask& task);
 #endif
 
+        // Returns the resume id.
+        std::string GetResumeId();
+
         // Called by the resume command. Loads the checkpoint manager with the resume id and returns the automatic checkpoint.
         std::optional<AppInstaller::Checkpoints::Checkpoint<AppInstaller::Checkpoints::AutomaticCheckpointData>> LoadCheckpoint(const std::string& resumeId);
 

src/AppInstallerCLICore/Resources.h

@@ -218,6 +218,7 @@ namespace AppInstaller::CLI::Resource
         WINGET_DEFINE_RESOURCE_STRINGID(HelpLinkPreamble);
         WINGET_DEFINE_RESOURCE_STRINGID(IdArgumentDescription);
         WINGET_DEFINE_RESOURCE_STRINGID(IgnoreLocalArchiveMalwareScanArgumentDescription);
+        WINGET_DEFINE_RESOURCE_STRINGID(IgnoreResumeLimitArgumentDescription);
         WINGET_DEFINE_RESOURCE_STRINGID(IgnoreWarningsArgumentDescription);
         WINGET_DEFINE_RESOURCE_STRINGID(ImportCommandLongDescription);
         WINGET_DEFINE_RESOURCE_STRINGID(ImportCommandReportDependencies);
@@ -425,6 +426,7 @@ namespace AppInstaller::CLI::Resource
         WINGET_DEFINE_RESOURCE_STRINGID(ResumeCommandShortDescription);
         WINGET_DEFINE_RESOURCE_STRINGID(ResumeIdArgumentDescription);
         WINGET_DEFINE_RESOURCE_STRINGID(ResumeIdNotFoundError);
+        WINGET_DEFINE_RESOURCE_STRINGID(ResumeLimitExceeded);
         WINGET_DEFINE_RESOURCE_STRINGID(ResumeStateDataNotFoundError);
         WINGET_DEFINE_RESOURCE_STRINGID(RetroArgumentDescription);
         WINGET_DEFINE_RESOURCE_STRINGID(SearchCommandLongDescription);

src/AppInstallerCLICore/Workflows/DependenciesFlow.cpp

@@ -228,7 +228,9 @@ namespace AppInstaller::CLI::Workflow
                 else
                 {
                     context.Reporter.Error() << Resource::String::RebootRequiredToEnableWindowsFeatureOverrideRequired << std::endl;
-                    AICLI_TERMINATE_CONTEXT(APPINSTALLER_CLI_ERROR_INSTALL_REBOOT_REQUIRED_TO_INSTALL);
+                    context.SetFlags(Execution::ContextFlag::RegisterResume);
+                    context.SetFlags(Execution::ContextFlag::RebootRequired);
+                    AICLI_TERMINATE_CONTEXT(APPINSTALLER_CLI_ERROR_INSTALL_REBOOT_REQUIRED_FOR_INSTALL);
                 }
             }
             else

src/AppInstallerCLICore/Workflows/InstallFlow.cpp

@@ -132,7 +132,7 @@ namespace AppInstaller::CLI::Workflow
                 case ExpectedReturnCodeEnum::RebootRequiredToFinish:
                     return ExpectedReturnCode(returnCode, APPINSTALLER_CLI_ERROR_INSTALL_REBOOT_REQUIRED_TO_FINISH, Resource::String::InstallFlowReturnCodeRebootRequiredToFinish);
                 case ExpectedReturnCodeEnum::RebootRequiredForInstall:
-                    return ExpectedReturnCode(returnCode, APPINSTALLER_CLI_ERROR_INSTALL_REBOOT_REQUIRED_TO_INSTALL, Resource::String::InstallFlowReturnCodeRebootRequiredForInstall);
+                    return ExpectedReturnCode(returnCode, APPINSTALLER_CLI_ERROR_INSTALL_REBOOT_REQUIRED_FOR_INSTALL, Resource::String::InstallFlowReturnCodeRebootRequiredForInstall);
                 case ExpectedReturnCodeEnum::RebootInitiated:
                     return ExpectedReturnCode(returnCode, APPINSTALLER_CLI_ERROR_INSTALL_REBOOT_INITIATED, Resource::String::InstallFlowReturnCodeRebootInitiated);
                 case ExpectedReturnCodeEnum::CancelledByUser:
@@ -484,8 +484,8 @@ namespace AppInstaller::CLI::Workflow
                     context.Reporter.Warn() << returnCode.Message << std::endl;
                     terminationHR = S_OK;
                     break;
-                case APPINSTALLER_CLI_ERROR_INSTALL_REBOOT_REQUIRED_TO_INSTALL:
-                    // REBOOT_REQUIRED_TO_INSTALL is treated as an error since installation has not yet completed.
+                case APPINSTALLER_CLI_ERROR_INSTALL_REBOOT_REQUIRED_FOR_INSTALL:
+                    // REBOOT_REQUIRED_FOR_INSTALL is treated as an error since installation has not yet completed.
                     context.SetFlags(ContextFlag::RebootRequired);
                     // TODO: Add separate workflow to handle restart registration for resume.
                     context.SetFlags(ContextFlag::RegisterResume);
@@ -599,7 +599,7 @@ namespace AppInstaller::CLI::Workflow
             Workflow::InstallDependencies <<
             Workflow::DownloadInstaller <<
             Workflow::InstallPackageInstaller <<
-            Workflow::InitiateRebootIfApplicable();
+            Workflow::RegisterStartupAfterReboot();
     }
 
     void EnsureSupportForInstall(Execution::Context& context)

src/AppInstallerCLICore/Workflows/ResumeFlow.cpp

@@ -3,6 +3,7 @@
 #include "pch.h"
 #include "ResumeFlow.h"
 #include "winget/Reboot.h"
+#include <AppInstallerRuntime.h>
 
 namespace AppInstaller::CLI::Workflow
 {
@@ -16,31 +17,22 @@ namespace AppInstaller::CLI::Workflow
         context.Checkpoint(m_checkpointName, m_contextData);
     }
 
-    void InitiateRebootIfApplicable::operator()(Execution::Context& context) const
+    void RegisterStartupAfterReboot::operator()(Execution::Context& context) const
     {
-        if (!Settings::ExperimentalFeature::IsEnabled(Settings::ExperimentalFeature::Feature::Reboot))
+        if (!Settings::ExperimentalFeature::IsEnabled(Settings::ExperimentalFeature::Feature::Resume) || 
+            !Settings::ExperimentalFeature::IsEnabled(Settings::ExperimentalFeature::Feature::Reboot))
         {
             return;
         }
 
-        if (!context.Args.Contains(Execution::Args::Type::AllowReboot))
+        if (WI_IsFlagSet(context.GetFlags(), Execution::ContextFlag::RegisterResume))
         {
-            AICLI_LOG(CLI, Info, << "No reboot flag found; skipping reboot flow.");
-            return;
-        }
-
-        if (WI_IsFlagSet(context.GetFlags(), Execution::ContextFlag::RebootRequired))
-        {
-            context.ClearFlags(Execution::ContextFlag::RebootRequired);
+            auto executablePath = AppInstaller::Runtime::GetPathTo(AppInstaller::Runtime::PathName::CLIExecutable);
 
-            if (Reboot::InitiateReboot())
-            {
-                context.Reporter.Warn() << Resource::String::InitiatingReboot << std::endl;
-            }
-            else
-            {
-                context.Reporter.Error() << Resource::String::FailedToInitiateReboot << std::endl;
-            }
+            // RunOnce registry value must start with the full path of the executable.
+            const auto& resumeId = context.GetResumeId();
+            std::string commandLine = executablePath.u8string() + " resume -g " + resumeId;
+            Reboot::WriteToRunOnceRegistry(resumeId, commandLine);
         }
     }
 }

src/AppInstallerCLICore/Workflows/ResumeFlow.h

@@ -23,14 +23,14 @@ namespace AppInstaller::CLI::Workflow
         std::vector<Execution::Data> m_contextData;
     };
 
-    // Initiates a reboot if applicable. This task always executes even if context terminates.
+    // Registers the resume command to execute upon reboot if applicable. This task always executes even if context terminates.
     // Required Args: None
     // Inputs: None
     // Outputs: None
-    struct InitiateRebootIfApplicable : public WorkflowTask
+    struct RegisterStartupAfterReboot : public WorkflowTask
     {
-        InitiateRebootIfApplicable() : WorkflowTask("InitiateRebootIfApplicable", /* executeAlways */true) {}
+        RegisterStartupAfterReboot() : WorkflowTask("RegisterStartupAfterReboot", /* executeAlways*/ true) {}
 
-        void operator()(Execution::Context& context) const override;
+        void operator()(Execution::Context & context) const override;
     };
 }

src/AppInstallerCLIE2ETests/Constants.cs

@@ -271,7 +271,7 @@ public class ErrorCode
             public const int ERROR_INSTALL_NO_NETWORK = unchecked((int)0x8A150107);
             public const int ERROR_INSTALL_CONTACT_SUPPORT = unchecked((int)0x8A150108);
             public const int ERROR_INSTALL_REBOOT_REQUIRED_TO_FINISH = unchecked((int)0x8A150109);
-            public const int ERROR_INSTALL_REBOOT_REQUIRED_TO_INSTALL = unchecked((int)0x8A15010A);
+            public const int ERROR_INSTALL_REBOOT_REQUIRED_FOR_INSTALL = unchecked((int)0x8A15010A);
             public const int ERROR_INSTALL_REBOOT_INITIATED = unchecked((int)0x8A15010B);
             public const int ERROR_INSTALL_CANCELLED_BY_USER = unchecked((int)0x8A15010C);
             public const int ERROR_INSTALL_ALREADY_INSTALLED = unchecked((int)0x8A15010D);

src/AppInstallerCLIE2ETests/Helpers/WinGetSettingsHelper.cs

@@ -44,6 +44,7 @@ public static void InitializeWingetSettings()
                         { "directMSI", false },
                         { "windowsFeature", false },
                         { "resume", false },
+                        { "reboot", false },
                     }
                 },
                 {
@@ -210,6 +211,7 @@ public static void InitializeAllFeatures(bool status)
             ConfigureFeature("directMSI", status);
             ConfigureFeature("windowsFeature", status);
             ConfigureFeature("resume", status);
+            ConfigureFeature("reboot", status);
         }
 
         private static JObject GetJsonSettingsObject(string objectName)