8 files changed, 173 insertions, 82 deletions

diff --git a/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h b/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
index c65ca86b..2633a9c2 100644
--- a/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
+++ b/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
@@ -293,6 +293,14 @@ enum BOOTSTRAPPER_EXECUTEPACKAGECOMPLETE_ACTION
     BOOTSTRAPPER_EXECUTEPACKAGECOMPLETE_ACTION_SUSPEND,
 };
 
+enum BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION
+{
+    BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION_NONE,
+    // Instructs the engine to stop processing the chain and restart.
+    // The engine will launch again after the machine is restarted.
+    BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION_RESTART,
+};
+
 enum BOOTSTRAPPER_EXECUTEPROCESSCANCEL_ACTION
 {
     // Instructs the engine to stop waiting for the process to exit.
@@ -697,11 +705,14 @@ struct BA_ONCOMMITMSITRANSACTIONCOMPLETE_ARGS
     DWORD cbSize;
     LPCWSTR wzTransactionId;
     HRESULT hrStatus;
+    BOOTSTRAPPER_APPLY_RESTART restart;
+    BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION recommendation;
 };
 
 struct BA_ONCOMMITMSITRANSACTIONCOMPLETE_RESULTS
 {
     DWORD cbSize;
+    BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION action;
 };
 
 struct BA_ONDETECTBEGIN_ARGS
@@ -1417,11 +1428,14 @@ struct BA_ONROLLBACKMSITRANSACTIONCOMPLETE_ARGS
     DWORD cbSize;
     LPCWSTR wzTransactionId;
     HRESULT hrStatus;
+    BOOTSTRAPPER_APPLY_RESTART restart;
+    BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION recommendation;
 };
 
 struct BA_ONROLLBACKMSITRANSACTIONCOMPLETE_RESULTS
 {
     DWORD cbSize;
+    BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION action;
 };
 
 struct BA_ONSETUPDATEBEGIN_ARGS
diff --git a/src/api/burn/WixToolset.Mba.Core/BootstrapperApplication.cs b/src/api/burn/WixToolset.Mba.Core/BootstrapperApplication.cs
index ecc99069..b23cb3b2 100644
--- a/src/api/burn/WixToolset.Mba.Core/BootstrapperApplication.cs
+++ b/src/api/burn/WixToolset.Mba.Core/BootstrapperApplication.cs
@@ -60,7 +60,7 @@ namespace WixToolset.Mba.Core
 
         /// <inheritdoc/>
         public event EventHandler<DetectCompatibleMsiPackageEventArgs> DetectCompatibleMsiPackage;
-        
+
         /// <inheritdoc/>
         public event EventHandler<DetectRelatedMsiPackageEventArgs> DetectRelatedMsiPackage;
 
@@ -1985,11 +1985,12 @@ namespace WixToolset.Mba.Core
             return args.HResult;
         }
 
-        int IBootstrapperApplication.OnCommitMsiTransactionComplete(string transactionId, int hrStatus)
+        int IBootstrapperApplication.OnCommitMsiTransactionComplete(string transactionId, int hrStatus, ApplyRestart restart, BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION recommendation, ref BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION pAction)
         {
-            CommitMsiTransactionCompleteEventArgs args = new CommitMsiTransactionCompleteEventArgs(transactionId, hrStatus);
+            CommitMsiTransactionCompleteEventArgs args = new CommitMsiTransactionCompleteEventArgs(transactionId, hrStatus, restart, recommendation, pAction);
             this.OnCommitMsiTransactionComplete(args);
 
+            pAction = args.Action;
             return args.HResult;
         }
 
@@ -2001,11 +2002,12 @@ namespace WixToolset.Mba.Core
             return args.HResult;
         }
 
-        int IBootstrapperApplication.OnRollbackMsiTransactionComplete(string transactionId, int hrStatus)
+        int IBootstrapperApplication.OnRollbackMsiTransactionComplete(string transactionId, int hrStatus, ApplyRestart restart, BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION recommendation, ref BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION pAction)
         {
-            RollbackMsiTransactionCompleteEventArgs args = new RollbackMsiTransactionCompleteEventArgs(transactionId, hrStatus);
+            RollbackMsiTransactionCompleteEventArgs args = new RollbackMsiTransactionCompleteEventArgs(transactionId, hrStatus, restart, recommendation, pAction);
             this.OnRollbackMsiTransactionComplete(args);
 
+            pAction = args.Action;
             return args.HResult;
         }
 
diff --git a/src/api/burn/WixToolset.Mba.Core/EventArgs.cs b/src/api/burn/WixToolset.Mba.Core/EventArgs.cs
index 48850fd7..9d18c4eb 100644
--- a/src/api/burn/WixToolset.Mba.Core/EventArgs.cs
+++ b/src/api/burn/WixToolset.Mba.Core/EventArgs.cs
@@ -568,7 +568,7 @@ namespace WixToolset.Mba.Core
     public class DetectPatchTargetEventArgs : CancellableHResultEventArgs
     {
         /// <summary>
-        /// 
+        ///
         /// </summary>
         /// <param name="packageId"></param>
         /// <param name="productCode"></param>
@@ -944,7 +944,7 @@ namespace WixToolset.Mba.Core
     public class PlanPatchTargetEventArgs : CancellableHResultEventArgs
     {
         /// <summary>
-        /// 
+        ///
         /// </summary>
         /// <param name="packageId"></param>
         /// <param name="productCode"></param>
@@ -1086,7 +1086,7 @@ namespace WixToolset.Mba.Core
     public class PlanPackageCompleteEventArgs : StatusEventArgs
     {
         /// <summary>
-        /// 
+        ///
         /// </summary>
         /// <param name="packageId"></param>
         /// <param name="hrStatus"></param>
@@ -2117,19 +2117,25 @@ namespace WixToolset.Mba.Core
     /// Additional arguments passed by the engine after committing an MSI transaction.
     /// </summary>
     [Serializable]
-    public class CommitMsiTransactionCompleteEventArgs : StatusEventArgs
+    public class CommitMsiTransactionCompleteEventArgs : ActionEventArgs<BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION>
     {
         /// <summary />
-        public CommitMsiTransactionCompleteEventArgs(string transactionId, int hrStatus)
-            : base(hrStatus)
+        public CommitMsiTransactionCompleteEventArgs(string transactionId, int hrStatus, ApplyRestart restart, BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION recommendation, BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION action)
+            : base(hrStatus, recommendation, action)
         {
             this.TransactionId = transactionId;
+            this.Restart = restart;
         }
 
         /// <summary>
         /// Gets the MSI transaction Id.
         /// </summary>
         public string TransactionId { get; private set; }
+
+        /// <summary>
+        /// Gets the package restart state after being applied.
+        /// </summary>
+        public ApplyRestart Restart { get; private set; }
     }
 
     /// <summary>
@@ -2154,19 +2160,25 @@ namespace WixToolset.Mba.Core
     /// Additional arguments passed by the engine after rolling back an MSI transaction.
     /// </summary>
     [Serializable]
-    public class RollbackMsiTransactionCompleteEventArgs : StatusEventArgs
+    public class RollbackMsiTransactionCompleteEventArgs : ActionEventArgs<BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION>
     {
         /// <summary />
-        public RollbackMsiTransactionCompleteEventArgs(string transactionId, int hrStatus)
-            : base(hrStatus)
+        public RollbackMsiTransactionCompleteEventArgs(string transactionId, int hrStatus, ApplyRestart restart, BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION recommendation, BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION action)
+            : base(hrStatus, recommendation, action)
         {
             this.TransactionId = transactionId;
+            this.Restart = restart;
         }
 
         /// <summary>
         /// Gets the MSI transaction Id.
         /// </summary>
         public string TransactionId { get; private set; }
+
+        /// <summary>
+        /// Gets the package restart state after being applied.
+        /// </summary>
+        public ApplyRestart Restart { get; private set; }
     }
 
     /// <summary>
@@ -2176,7 +2188,7 @@ namespace WixToolset.Mba.Core
     public class PauseAutomaticUpdatesBeginEventArgs : HResultEventArgs
     {
         /// <summary>
-        /// 
+        ///
         /// </summary>
         public PauseAutomaticUpdatesBeginEventArgs()
         {
@@ -2190,7 +2202,7 @@ namespace WixToolset.Mba.Core
     public class PauseAutomaticUpdatesCompleteEventArgs : StatusEventArgs
     {
         /// <summary>
-        /// 
+        ///
         /// </summary>
         /// <param name="hrStatus"></param>
         public PauseAutomaticUpdatesCompleteEventArgs(int hrStatus)
@@ -2206,7 +2218,7 @@ namespace WixToolset.Mba.Core
     public class SystemRestorePointBeginEventArgs : HResultEventArgs
     {
         /// <summary>
-        /// 
+        ///
         /// </summary>
         public SystemRestorePointBeginEventArgs()
         {
@@ -2220,7 +2232,7 @@ namespace WixToolset.Mba.Core
     public class SystemRestorePointCompleteEventArgs : StatusEventArgs
     {
         /// <summary>
-        /// 
+        ///
         /// </summary>
         /// <param name="hrStatus"></param>
         public SystemRestorePointCompleteEventArgs(int hrStatus)
diff --git a/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplication.cs b/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplication.cs
index daa95e17..7bfabeb3 100644
--- a/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplication.cs
+++ b/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplication.cs
@@ -933,12 +933,18 @@ namespace WixToolset.Mba.Core
         /// </summary>
         /// <param name="wzTransactionId"></param>
         /// <param name="hrStatus"></param>
+        /// <param name="restart"></param>
+        /// <param name="recommendation"></param>
+        /// <param name="pAction"></param>
         /// <returns></returns>
         [PreserveSig]
         [return: MarshalAs(UnmanagedType.I4)]
         int OnCommitMsiTransactionComplete(
             [MarshalAs(UnmanagedType.LPWStr)] string wzTransactionId,
-            int hrStatus
+            int hrStatus,
+            [MarshalAs(UnmanagedType.U4)] ApplyRestart restart,
+            [MarshalAs(UnmanagedType.I4)] BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION recommendation,
+            [MarshalAs(UnmanagedType.I4)] ref BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION pAction
             );
 
         /// <summary>
@@ -957,12 +963,18 @@ namespace WixToolset.Mba.Core
         /// </summary>
         /// <param name="wzTransactionId"></param>
         /// <param name="hrStatus"></param>
+        /// <param name="restart"></param>
+        /// <param name="recommendation"></param>
+        /// <param name="pAction"></param>
         /// <returns></returns>
         [PreserveSig]
         [return: MarshalAs(UnmanagedType.I4)]
         int OnRollbackMsiTransactionComplete(
             [MarshalAs(UnmanagedType.LPWStr)] string wzTransactionId,
-            int hrStatus
+            int hrStatus,
+            [MarshalAs(UnmanagedType.U4)] ApplyRestart restart,
+            [MarshalAs(UnmanagedType.I4)] BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION recommendation,
+            [MarshalAs(UnmanagedType.I4)] ref BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION pAction
             );
 
         /// <summary>
@@ -1199,27 +1211,27 @@ namespace WixToolset.Mba.Core
     public enum Display
     {
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Unknown,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Embedded,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         None,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Passive,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Full,
     }
@@ -1331,27 +1343,27 @@ namespace WixToolset.Mba.Core
     public enum Restart
     {
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Unknown,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Never,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Prompt,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Automatic,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Always,
     }
@@ -1384,47 +1396,47 @@ namespace WixToolset.Mba.Core
     public enum Result
     {
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Error = -1,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         None,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Ok,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Cancel,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Abort,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Retry,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Ignore,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Yes,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         No,
 
@@ -1434,17 +1446,17 @@ namespace WixToolset.Mba.Core
         Close,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Help,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         TryAgain,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Continue,
     }
@@ -1455,7 +1467,7 @@ namespace WixToolset.Mba.Core
     public enum ResumeType
     {
         /// <summary>
-        /// 
+        ///
         /// </summary>
         None,
 
@@ -1546,7 +1558,7 @@ namespace WixToolset.Mba.Core
     public enum RelatedOperation
     {
         /// <summary>
-        /// 
+        ///
         /// </summary>
         None,
 
@@ -1685,47 +1697,47 @@ namespace WixToolset.Mba.Core
     public enum RelationType
     {
         /// <summary>
-        /// 
+        ///
         /// </summary>
         None,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Detect,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Upgrade,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Addon,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Patch,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         DependentAddon,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         DependentPatch,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Update,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         ChainPackage,
     }
@@ -1736,37 +1748,37 @@ namespace WixToolset.Mba.Core
     public enum RelatedBundlePlanType
     {
         /// <summary>
-        /// 
+        ///
         /// </summary>
         None,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Downgrade,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Upgrade,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Addon,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         Patch,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         DependentAddon,
 
         /// <summary>
-        /// 
+        ///
         /// </summary>
         DependentPatch,
     }
@@ -1804,7 +1816,7 @@ namespace WixToolset.Mba.Core
     public enum BOOTSTRAPPER_APPLYCOMPLETE_ACTION
     {
         /// <summary>
-        /// 
+        ///
         /// </summary>
         None,
 
@@ -1843,7 +1855,7 @@ namespace WixToolset.Mba.Core
     public enum BOOTSTRAPPER_CACHEACQUIRECOMPLETE_ACTION
     {
         /// <summary>
-        /// 
+        ///
         /// </summary>
         None,
 
@@ -1860,7 +1872,7 @@ namespace WixToolset.Mba.Core
     public enum BOOTSTRAPPER_CACHEPACKAGECOMPLETE_ACTION
     {
         /// <summary>
-        /// 
+        ///
         /// </summary>
         None,
 
@@ -1883,7 +1895,7 @@ namespace WixToolset.Mba.Core
     public enum BOOTSTRAPPER_CACHEPACKAGENONVITALVALIDATIONFAILURE_ACTION
     {
         /// <summary>
-        /// 
+        ///
         /// </summary>
         None,
 
@@ -1900,7 +1912,7 @@ namespace WixToolset.Mba.Core
     public enum BOOTSTRAPPER_CACHEVERIFYCOMPLETE_ACTION
     {
         /// <summary>
-        /// 
+        ///
         /// </summary>
         None,
 
@@ -1921,7 +1933,7 @@ namespace WixToolset.Mba.Core
     public enum BOOTSTRAPPER_EXECUTEPACKAGECOMPLETE_ACTION
     {
         /// <summary>
-        /// 
+        ///
         /// </summary>
         None,
 
@@ -1950,6 +1962,23 @@ namespace WixToolset.Mba.Core
     }
 
     /// <summary>
+    /// The available actions for <see cref="IDefaultBootstrapperApplication.CommitMsiTransactionComplete"/> and <see cref="IDefaultBootstrapperApplication.RollbackMsiTransactionComplete"/>.
+    /// </summary>
+    public enum BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION
+    {
+        /// <summary>
+        ///
+        /// </summary>
+        None,
+
+        /// <summary>
+        /// Instructs the engine to stop processing the chain and restart.
+        /// The engine will launch again after the machine is restarted.
+        /// </summary>
+        Restart,
+    };
+
+    /// <summary>
     /// The available actions for <see cref="IDefaultBootstrapperApplication.ExecuteProcessCancel"/>.
     /// </summary>
     public enum BOOTSTRAPPER_EXECUTEPROCESSCANCEL_ACTION
@@ -2017,7 +2046,7 @@ namespace WixToolset.Mba.Core
     public enum BOOTSTRAPPER_SHUTDOWN_ACTION
     {
         /// <summary>
-        /// 
+        ///
         /// </summary>
         None,
 
diff --git a/src/api/burn/balutil/inc/BalBaseBAFunctions.h b/src/api/burn/balutil/inc/BalBaseBAFunctions.h
index 9a8ac87e..6ad109c3 100644
--- a/src/api/burn/balutil/inc/BalBaseBAFunctions.h
+++ b/src/api/burn/balutil/inc/BalBaseBAFunctions.h
@@ -710,7 +710,10 @@ public: // IBootstrapperApplication
 
     virtual STDMETHODIMP OnCommitMsiTransactionComplete(
         __in_z LPCWSTR /*wzTransactionId*/,
-        __in HRESULT /*hrStatus*/
+        __in HRESULT /*hrStatus*/,
+        __in BOOTSTRAPPER_APPLY_RESTART /*restart*/,
+        __in BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION /*recommendation*/,
+        __inout BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION* /*pAction*/
         )
     {
         return S_OK;
@@ -725,7 +728,10 @@ public: // IBootstrapperApplication
 
     virtual STDMETHODIMP OnRollbackMsiTransactionComplete(
         __in_z LPCWSTR /*wzTransactionId*/,
-        __in HRESULT /*hrStatus*/
+        __in HRESULT /*hrStatus*/,
+        __in BOOTSTRAPPER_APPLY_RESTART /*restart*/,
+        __in BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION /*recommendation*/,
+        __inout BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION* /*pAction*/
         )
     {
         return S_OK;
diff --git a/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h b/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h
index 58cc0673..c8b80d13 100644
--- a/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h
+++ b/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h
@@ -197,7 +197,7 @@ public: // IBootstrapperApplication
         __in LPCWSTR /*wzVersion*/,
         __in BOOTSTRAPPER_RELATED_OPERATION /*operation*/,
         __inout BOOL* pfCancel
-        ) 
+        )
     {
         *pfCancel |= CheckCanceled();
         return S_OK;
@@ -902,10 +902,21 @@ public: // IBootstrapperApplication
 
     virtual STDMETHODIMP OnCommitMsiTransactionComplete(
         __in_z LPCWSTR /*wzTransactionId*/,
-        __in HRESULT /*hrStatus*/
+        __in HRESULT /*hrStatus*/,
+        __in BOOTSTRAPPER_APPLY_RESTART /*restart*/,
+        __in BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION /*recommendation*/,
+        __inout BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION* /*pAction*/
         )
     {
-        return S_OK;
+        HRESULT hr = S_OK;
+
+        if (CheckCanceled())
+        {
+            ExitFunction1(hr = HRESULT_FROM_WIN32(ERROR_INSTALL_USEREXIT));
+        }
+
+    LExit:
+        return hr;
     }
 
     virtual STDMETHODIMP OnRollbackMsiTransactionBegin(
@@ -917,10 +928,21 @@ public: // IBootstrapperApplication
 
     virtual STDMETHODIMP OnRollbackMsiTransactionComplete(
         __in_z LPCWSTR /*wzTransactionId*/,
-        __in HRESULT /*hrStatus*/
+        __in HRESULT /*hrStatus*/,
+        __in BOOTSTRAPPER_APPLY_RESTART /*restart*/,
+        __in BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION /*recommendation*/,
+        __inout BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION* /*pAction*/
         )
     {
-        return S_OK;
+        HRESULT hr = S_OK;
+
+        if (CheckCanceled())
+        {
+            ExitFunction1(hr = HRESULT_FROM_WIN32(ERROR_INSTALL_USEREXIT));
+        }
+
+    LExit:
+        return hr;
     }
 
     virtual STDMETHODIMP OnPauseAutomaticUpdatesBegin(
diff --git a/src/api/burn/balutil/inc/BalBaseBootstrapperApplicationProc.h b/src/api/burn/balutil/inc/BalBaseBootstrapperApplicationProc.h
index 62cb85bc..34739ba3 100644
--- a/src/api/burn/balutil/inc/BalBaseBootstrapperApplicationProc.h
+++ b/src/api/burn/balutil/inc/BalBaseBootstrapperApplicationProc.h
@@ -579,10 +579,10 @@ static HRESULT BalBaseBAProcOnCommitMsiTransactionBegin(
 static HRESULT BalBaseBAProcOnCommitMsiTransactionComplete(
     __in IBootstrapperApplication* pBA,
     __in BA_ONCOMMITMSITRANSACTIONCOMPLETE_ARGS* pArgs,
-    __inout BA_ONCOMMITMSITRANSACTIONCOMPLETE_RESULTS* /*pResults*/
+    __inout BA_ONCOMMITMSITRANSACTIONCOMPLETE_RESULTS* pResults
     )
 {
-    return pBA->OnCommitMsiTransactionComplete(pArgs->wzTransactionId, pArgs->hrStatus);
+    return pBA->OnCommitMsiTransactionComplete(pArgs->wzTransactionId, pArgs->hrStatus, pArgs->restart, pArgs->recommendation, &pResults->action);
 }
 
 static HRESULT BalBaseBAProcOnRollbackMsiTransactionBegin(
@@ -597,10 +597,10 @@ static HRESULT BalBaseBAProcOnRollbackMsiTransactionBegin(
 static HRESULT BalBaseBAProcOnRollbackMsiTransactionComplete(
     __in IBootstrapperApplication* pBA,
     __in BA_ONROLLBACKMSITRANSACTIONCOMPLETE_ARGS* pArgs,
-    __inout BA_ONROLLBACKMSITRANSACTIONCOMPLETE_RESULTS* /*pResults*/
+    __inout BA_ONROLLBACKMSITRANSACTIONCOMPLETE_RESULTS* pResults
     )
 {
-    return pBA->OnRollbackMsiTransactionComplete(pArgs->wzTransactionId, pArgs->hrStatus);
+    return pBA->OnRollbackMsiTransactionComplete(pArgs->wzTransactionId, pArgs->hrStatus, pArgs->restart, pArgs->recommendation, &pResults->action);
 }
 
 static HRESULT BalBaseBAProcOnPauseAutomaticUpdatesBegin(
@@ -780,7 +780,7 @@ static HRESULT WINAPI BalBaseBootstrapperApplicationProc(
 {
     IBootstrapperApplication* pBA = reinterpret_cast<IBootstrapperApplication*>(pvContext);
     HRESULT hr = pBA->BAProc(message, pvArgs, pvResults, pvContext);
-    
+
     if (E_NOTIMPL == hr)
     {
         switch (message)
diff --git a/src/api/burn/balutil/inc/IBootstrapperApplication.h b/src/api/burn/balutil/inc/IBootstrapperApplication.h
index fba919ea..faa211e4 100644
--- a/src/api/burn/balutil/inc/IBootstrapperApplication.h
+++ b/src/api/burn/balutil/inc/IBootstrapperApplication.h
@@ -574,7 +574,7 @@ DECLARE_INTERFACE_IID_(IBootstrapperApplication, IUnknown, "53C31D56-49C0-426B-A
         ) = 0;
 
     // OnLaunchApprovedExeBegin - called before trying to launch the preapproved executable.
-    // 
+    //
     STDMETHOD(OnLaunchApprovedExeBegin)(
         __inout BOOL* pfCancel
         ) = 0;
@@ -603,7 +603,10 @@ DECLARE_INTERFACE_IID_(IBootstrapperApplication, IUnknown, "53C31D56-49C0-426B-A
 
     STDMETHOD(OnCommitMsiTransactionComplete)(
         __in_z LPCWSTR wzTransactionId,
-        __in HRESULT hrStatus
+        __in HRESULT hrStatus,
+        __in BOOTSTRAPPER_APPLY_RESTART restart,
+        __in BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION recommendation,
+        __inout BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION* pAction
         ) = 0;
 
     STDMETHOD(OnRollbackMsiTransactionBegin)(
@@ -612,7 +615,10 @@ DECLARE_INTERFACE_IID_(IBootstrapperApplication, IUnknown, "53C31D56-49C0-426B-A
 
     STDMETHOD(OnRollbackMsiTransactionComplete)(
         __in_z LPCWSTR wzTransactionId,
-        __in HRESULT hrStatus
+        __in HRESULT hrStatus,
+        __in BOOTSTRAPPER_APPLY_RESTART restart,
+        __in BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION recommendation,
+        __inout BOOTSTRAPPER_EXECUTEMSITRANSACTIONCOMPLETE_ACTION* pAction
         ) = 0;
 
     STDMETHOD(OnPauseAutomaticUpdatesBegin)(