diff --git a/frameworks/core/interfaces/native/common/api_impl.cpp b/frameworks/core/interfaces/native/common/api_impl.cpp
index c4717ebb574a7d1c9ab451ec4c81838a40fda73d..7a570f9dbb69ffa53b7669f3075a06c37d600e52 100644
--- a/frameworks/core/interfaces/native/common/api_impl.cpp
+++ b/frameworks/core/interfaces/native/common/api_impl.cpp
@@ -343,6 +343,18 @@ void* GetAttachNodePtr(Ark_NodeHandle nodePtr)
     return frameNode->GetExternalData();
 }
 
+static void SetVMContext(Ark_VMContext vmContext) {
+    auto pipeline = PipelineContext::GetCurrentContextSafely();
+    CHECK_NULL_VOID(pipeline);
+    auto rootNode = pipeline->GetRootElement();
+    CHECK_NULL_VOID(rootNode);
+    auto rootPtr = reinterpret_cast<Ark_NodeHandle>(Referenced::RawPtr(rootNode));
+    CHECK_NULL_VOID(rootPtr);
+    auto companionNode = GeneratedApiImpl::GetCompanion(rootPtr);
+    CHECK_NULL_VOID(companionNode);
+    companionNode->SetVMContext(vmContext);
+}
+
 Ark_Int32 MeasureLayoutAndDraw(Ark_VMContext vmContext, Ark_NodeHandle rootPtr)
 {
     auto* root = reinterpret_cast<FrameNode*>(rootPtr);
@@ -365,9 +377,7 @@ Ark_Int32 MeasureLayoutAndDraw(Ark_VMContext vmContext, Ark_NodeHandle rootPtr)
     Ark_Float32 drawData[] = { 0, 0, 0, 0 };
     DrawNode(vmContext, stagePtr, &drawData[0]);
 
-    if (auto companionNode = GetCompanion(stagePtr); companionNode) {
-        companionNode->SetVMContext(vmContext);
-    }
+    SetVMContext(vmContext);
     return 0;
 }