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/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,45 @@ 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);
-        }
+        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();
+            }
+
+            context.ClearFlags(Execution::ContextFlag::RebootRequired);
+
+            if (!Reboot::InitiateReboot())
+            {
+                context.Reporter.Error() << Resource::String::FailedToInitiateReboot << std::endl;
+            }
+            else if (WI_IsFlagSet(context.GetFlags(), Execution::ContextFlag::RegisterResume))
+            {
+                context.ClearFlags(Execution::ContextFlag::RegisterResume);
 
-        if (context.Args.Contains(Execution::Args::Type::Wait))
+                // For Windows Error Reporting to restart this process, the process must be rebooted while it is still running.
+                // Workaround is to have the program sleep while the reboot process is kicked off. Only applies to resume.
+                Sleep(5000);
+            }
+        }
+        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

@@ -99,6 +99,18 @@ 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 (resumeCount >= Settings::User().Get<Settings::Setting::MaxResumes>())
+        {
+            context.Reporter.Error() << Resource::String::ResumeLimitExceeded(Utility::LocIndView{ resumeCountString }) << 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/ExecutionContext.cpp

@@ -255,9 +255,9 @@ namespace AppInstaller::CLI::Execution
 
     Context::~Context()
     {
-        if (Settings::ExperimentalFeature::IsEnabled(ExperimentalFeature::Feature::Resume) && !IsTerminated())
+        if (Settings::ExperimentalFeature::IsEnabled(ExperimentalFeature::Feature::Resume))
         {
-            if (m_checkpointManager)
+            if (m_checkpointManager && (!IsTerminated() || GetTerminationHR() != APPINSTALLER_CLI_ERROR_INSTALL_REBOOT_REQUIRED_FOR_INSTALL))
             {
                 m_checkpointManager->CleanUpDatabase();
             }
@@ -427,6 +427,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);

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

@@ -187,6 +187,7 @@ namespace AppInstaller::CLI::Resource
         WINGET_DEFINE_RESOURCE_STRINGID(FailedToEnableWindowsFeatureOverrideRequired);
         WINGET_DEFINE_RESOURCE_STRINGID(FailedToInitiateReboot);
         WINGET_DEFINE_RESOURCE_STRINGID(FailedToRefreshPathWarning);
+        WINGET_DEFINE_RESOURCE_STRINGID(FailedToRegisterReboot);
         WINGET_DEFINE_RESOURCE_STRINGID(FeatureDisabledByAdminSettingMessage);
         WINGET_DEFINE_RESOURCE_STRINGID(FeatureDisabledMessage);
         WINGET_DEFINE_RESOURCE_STRINGID(FeaturesCommandLongDescription);
@@ -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::RegisterForReboot();
     }
 
     void EnsureSupportForInstall(Execution::Context& context)

src/AppInstallerCLICore/Workflows/ResumeFlow.cpp

@@ -16,30 +16,26 @@ namespace AppInstaller::CLI::Workflow
         context.Checkpoint(m_checkpointName, m_contextData);
     }
 
-    void InitiateRebootIfApplicable::operator()(Execution::Context& context) const
+    void RegisterForReboot::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);
+            std::string commandLineArg = "resume -g " + context.GetResumeId();
 
-            if (Reboot::InitiateReboot())
+            if (!Reboot::RegisterApplicationForReboot(commandLineArg))
             {
-                context.Reporter.Warn() << Resource::String::InitiatingReboot << std::endl;
-            }
-            else
-            {
-                context.Reporter.Error() << Resource::String::FailedToInitiateReboot << std::endl;
+                context.Reporter.Error() << Resource::String::FailedToRegisterReboot << std::endl;
+
+                if (WI_IsFlagSet(context.GetFlags(), Execution::ContextFlag::RebootRequired))
+                {
+                    context.ClearFlags(Execution::ContextFlag::RebootRequired);
+                }
             }
         }
     }

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 RegisterForReboot : public WorkflowTask
     {
-        InitiateRebootIfApplicable() : WorkflowTask("InitiateRebootIfApplicable", /* executeAlways */true) {}
+        RegisterForReboot() : WorkflowTask("RegisterForReboot", /* 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)

src/AppInstallerCLIE2ETests/ResumeCommand.cs

@@ -67,24 +67,46 @@ public void ResumeIdNotFound()
             Assert.AreEqual(Constants.ErrorCode.ERROR_RESUME_ID_NOT_FOUND, resumeResult.ExitCode);
         }
 
-        /// <summary>
-        /// Verifies that a checkpoint record persists after a failed install.
+        /// <summary>
+        /// Test install a package that returns REBOOT_REQUIRED_TO_FINISH and verify that the checkpoint database is properly cleaned up.
+        /// </summary>
+        [Test]
+        public void InstallRequiresRebootToFinish()
+        {
+            var checkpointsDir = TestCommon.GetCheckpointsDirectory();
+            int initialCheckpointsCount = Directory.Exists(checkpointsDir) ? Directory.GetDirectories(checkpointsDir).Length : 0;
+
+            var installDir = TestCommon.GetRandomTestDir();
+            var result = TestCommon.RunAICLICommand("install", $"TestRebootRequired --custom \"/ExitCode 9\" -l {installDir}");
+
+            // REBOOT_REQUIRED_TO_FINISH is treated as a success.
+            Assert.AreEqual(Constants.ErrorCode.S_OK, result.ExitCode);
+            Assert.True(result.StdOut.Contains("Restart your PC to finish installation."));
+            Assert.True(TestCommon.VerifyTestExeInstalledAndCleanup(installDir));
+
+            int actualCheckpointsCount = Directory.GetDirectories(checkpointsDir).Length;
+
+            // Checkpoint database should be cleaned up since resume is not needed to complete installation.
+            Assert.AreEqual(initialCheckpointsCount, actualCheckpointsCount);
+        }
+
+        /// <summary>
+        /// Test install a package that returns REBOOT_REQUIRED_FOR_INSTALL and verify that resume command can be called successfully.
         /// </summary>
         [Test]
-        public void ResumeRecordPreserved()
-        {
+        public void InstallRequiresRebootForInstall()
+        {
             var checkpointsDir = TestCommon.GetCheckpointsDirectory();
 
             int initialCheckpointsCount = Directory.Exists(checkpointsDir) ? Directory.GetDirectories(checkpointsDir).Length : 0;
 
-            // TODO: Refine test case to more accurately reflect usage once resume is fully implemented.
-            var result = TestCommon.RunAICLICommand("install", $"AppInstallerTest.TestZipInvalidRelativePath");
-            Assert.AreNotEqual(Constants.ErrorCode.S_OK, result.ExitCode);
-            Assert.True(result.StdOut.Contains("Invalid relative file path to the nested installer; path points to a location outside of the install directory"));
+            var installDir = TestCommon.GetRandomTestDir();
+            var result = TestCommon.RunAICLICommand("install", $"TestRebootRequired --custom \"/ExitCode 10\" -l {installDir}");
+            Assert.AreEqual(Constants.ErrorCode.ERROR_INSTALL_REBOOT_REQUIRED_FOR_INSTALL, result.ExitCode);
+            Assert.True(result.StdOut.Contains("Your PC will restart to finish installation."));
+            Assert.True(TestCommon.VerifyTestExeInstalledAndCleanup(installDir));
 
             int actualCheckpointsCount = Directory.GetDirectories(checkpointsDir).Length;
-
-            // One new checkpoint record should be created after running the install command.
             Assert.AreEqual(initialCheckpointsCount + 1, actualCheckpointsCount);
 
             var checkpointsDirectoryInfo = new DirectoryInfo(checkpointsDir);
@@ -95,8 +117,9 @@ public void ResumeRecordPreserved()
 
             // Resume output should be the same as the install result.
             var resumeResult = TestCommon.RunAICLICommand("resume", $"-g {checkpoint.Name}");
-            Assert.AreNotEqual(Constants.ErrorCode.S_OK, resumeResult.ExitCode);
-            Assert.True(resumeResult.StdOut.Contains("Invalid relative file path to the nested installer; path points to a location outside of the install directory"));
+            Assert.AreEqual(Constants.ErrorCode.ERROR_INSTALL_REBOOT_REQUIRED_FOR_INSTALL, resumeResult.ExitCode);
+            Assert.True(resumeResult.StdOut.Contains("Your PC will restart to finish installation."));
+            Assert.True(TestCommon.VerifyTestExeInstalledAndCleanup(installDir));
         }
     }
 }