1 files changed, 7 insertions, 5 deletions

diff --git a/src/api/burn/balutil/inc/BAFunctions.h b/src/api/burn/balutil/inc/BAFunctions.h
index 4a0b8599..7908a74c 100644
--- a/src/api/burn/balutil/inc/BAFunctions.h
+++ b/src/api/burn/balutil/inc/BAFunctions.h
@@ -1,6 +1,7 @@
 #pragma once
 // Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information.
 
+#include "IBootstrapperEngine.h"
 
 #ifdef __cplusplus
 extern "C" {
@@ -9,12 +10,14 @@ extern "C" {
 // The first 1024 messages are reserved so that the BA messages have the same value here.
 enum BA_FUNCTIONS_MESSAGE
 {
+    BA_FUNCTIONS_MESSAGE_ONCREATE = BOOTSTRAPPER_APPLICATION_MESSAGE_ONCREATE,
+    BA_FUNCTIONS_MESSAGE_ONDESTROY = BOOTSTRAPPER_APPLICATION_MESSAGE_ONDESTROY,
+    BA_FUNCTIONS_MESSAGE_ONSTARTUP = BOOTSTRAPPER_APPLICATION_MESSAGE_ONSTARTUP,
+    BA_FUNCTIONS_MESSAGE_ONSHUTDOWN = BOOTSTRAPPER_APPLICATION_MESSAGE_ONSHUTDOWN,
     BA_FUNCTIONS_MESSAGE_ONDETECTBEGIN = BOOTSTRAPPER_APPLICATION_MESSAGE_ONDETECTBEGIN,
     BA_FUNCTIONS_MESSAGE_ONDETECTCOMPLETE = BOOTSTRAPPER_APPLICATION_MESSAGE_ONDETECTCOMPLETE,
     BA_FUNCTIONS_MESSAGE_ONPLANBEGIN = BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANBEGIN,
     BA_FUNCTIONS_MESSAGE_ONPLANCOMPLETE = BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANCOMPLETE,
-    BA_FUNCTIONS_MESSAGE_ONSTARTUP = BOOTSTRAPPER_APPLICATION_MESSAGE_ONSTARTUP,
-    BA_FUNCTIONS_MESSAGE_ONSHUTDOWN = BOOTSTRAPPER_APPLICATION_MESSAGE_ONSHUTDOWN,
     BA_FUNCTIONS_MESSAGE_ONDETECTFORWARDCOMPATIBLEBUNDLE = BOOTSTRAPPER_APPLICATION_MESSAGE_ONDETECTFORWARDCOMPATIBLEBUNDLE,
     BA_FUNCTIONS_MESSAGE_ONDETECTUPDATEBEGIN = BOOTSTRAPPER_APPLICATION_MESSAGE_ONDETECTUPDATEBEGIN,
     BA_FUNCTIONS_MESSAGE_ONDETECTUPDATE = BOOTSTRAPPER_APPLICATION_MESSAGE_ONDETECTUPDATE,
@@ -81,8 +84,6 @@ enum BA_FUNCTIONS_MESSAGE
     BA_FUNCTIONS_MESSAGE_ONCACHEPAYLOADEXTRACTCOMPLETE = BOOTSTRAPPER_APPLICATION_MESSAGE_ONCACHEPAYLOADEXTRACTCOMPLETE,
     BA_FUNCTIONS_MESSAGE_ONCACHEPAYLOADEXTRACTPROGRESS = BOOTSTRAPPER_APPLICATION_MESSAGE_ONCACHEPAYLOADEXTRACTPROGRESS,
     BA_FUNCTIONS_MESSAGE_ONPLANROLLBACKBOUNDARY = BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANROLLBACKBOUNDARY,
-    BA_FUNCTIONS_MESSAGE_ONSETUPDATEBEGIN = BOOTSTRAPPER_APPLICATION_MESSAGE_ONSETUPDATEBEGIN,
-    BA_FUNCTIONS_MESSAGE_ONSETUPDATECOMPLETE = BOOTSTRAPPER_APPLICATION_MESSAGE_ONSETUPDATECOMPLETE,
     BA_FUNCTIONS_MESSAGE_ONDETECTCOMPATIBLEMSIPACKAGE = BOOTSTRAPPER_APPLICATION_MESSAGE_ONDETECTCOMPATIBLEMSIPACKAGE,
     BA_FUNCTIONS_MESSAGE_ONPLANCOMPATIBLEMSIPACKAGEBEGIN = BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANCOMPATIBLEMSIPACKAGEBEGIN,
     BA_FUNCTIONS_MESSAGE_ONPLANCOMPATIBLEMSIPACKAGECOMPLETE = BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANCOMPATIBLEMSIPACKAGECOMPLETE,
@@ -117,7 +118,8 @@ struct BA_FUNCTIONS_CREATE_ARGS
 {
     DWORD cbSize;
     DWORD64 qwBAFunctionsAPIVersion;
-    BOOTSTRAPPER_CREATE_ARGS* pBootstrapperCreateArgs;
+    IBootstrapperEngine* pEngine;
+    BOOTSTRAPPER_COMMAND* pCommand;
 };
 
 struct BA_FUNCTIONS_CREATE_RESULTS