diff --git a/TEST_MAPPING b/TEST_MAPPING
index 8cb01d7..5dd060a 100644
--- a/TEST_MAPPING
+++ b/TEST_MAPPING
@@ -49,6 +49,12 @@
     },
     {
       "name": "vm_accessor_test"
+    },
+    {
+      "name": "avf_backcompat_tests"
+    },
+    {
+      "name": "old_images_avf_test"
     }
   ],
   "avf-postsubmit": [
diff --git a/android/TerminalApp/Android.bp b/android/TerminalApp/Android.bp
index 545ba0f..2bac412 100644
--- a/android/TerminalApp/Android.bp
+++ b/android/TerminalApp/Android.bp
@@ -16,6 +16,7 @@
         "androidx-constraintlayout_constraintlayout",
         "androidx.window_window",
         "apache-commons-compress",
+        "avf_aconfig_flags_java",
         "com.google.android.material_material",
         "debian-service-grpclib-lite",
         "gson",
@@ -40,6 +41,7 @@
         //optimize: true,
         proguard_flags_files: ["proguard.flags"],
         shrink_resources: true,
+        keep_runtime_invisible_annotations: true,
     },
     apex_available: [
         "com.android.virt",
diff --git a/android/TerminalApp/AndroidManifest.xml b/android/TerminalApp/AndroidManifest.xml
index 53fdafc..c11b1a0 100644
--- a/android/TerminalApp/AndroidManifest.xml
+++ b/android/TerminalApp/AndroidManifest.xml
@@ -47,6 +47,11 @@
                 <category android:name="android.intent.category.DEFAULT" />
             </intent-filter>
         </activity>
+        <activity android:name=".DisplayActivity"
+            android:screenOrientation="landscape"
+            android:resizeableActivity="false"
+            android:theme="@style/FullscreenTheme"
+            android:configChanges="orientation|screenSize|keyboard|keyboardHidden|navigation|uiMode|screenLayout|smallestScreenSize" />
         <activity android:name=".SettingsActivity"
             android:label="@string/action_settings" />
         <activity android:name=".SettingsDiskResizeActivity"
diff --git a/android/TerminalApp/java/com/android/virtualization/terminal/DisplayActivity.kt b/android/TerminalApp/java/com/android/virtualization/terminal/DisplayActivity.kt
new file mode 100644
index 0000000..290cf5a
--- /dev/null
+++ b/android/TerminalApp/java/com/android/virtualization/terminal/DisplayActivity.kt
@@ -0,0 +1,65 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.virtualization.terminal
+
+import android.os.Bundle
+import android.system.virtualmachine.VirtualMachineManager
+import android.view.SurfaceView
+import android.view.WindowInsets
+import android.view.WindowInsetsController
+
+class DisplayActivity : BaseActivity() {
+    private lateinit var displayProvider: DisplayProvider
+
+    override fun onCreate(savedInstanceState: Bundle?) {
+        super.onCreate(savedInstanceState)
+        setContentView(R.layout.activity_display)
+        val mainView = findViewById<SurfaceView>(R.id.surface_view)
+        val cursorView = findViewById<SurfaceView>(R.id.cursor_surface_view)
+        makeFullscreen()
+        // Connect the views to the VM
+        displayProvider = DisplayProvider(mainView, cursorView)
+        val vmm =
+            applicationContext.getSystemService<VirtualMachineManager>(
+                VirtualMachineManager::class.java
+            )
+        val debianVm = vmm.get("debian")
+        if (debianVm != null) {
+            InputForwarder(
+                this,
+                debianVm,
+                findViewById(R.id.background_touch_view),
+                findViewById(R.id.surface_view),
+                findViewById(R.id.surface_view),
+            )
+        }
+    }
+
+    override fun onPause() {
+        super.onPause()
+        displayProvider.notifyDisplayIsGoingToInvisible()
+    }
+
+    private fun makeFullscreen() {
+        val w = window
+        w.setDecorFitsSystemWindows(false)
+        val insetsCtrl = w.insetsController
+        insetsCtrl?.hide(WindowInsets.Type.systemBars())
+        insetsCtrl?.setSystemBarsBehavior(
+            WindowInsetsController.BEHAVIOR_SHOW_TRANSIENT_BARS_BY_SWIPE
+        )
+    }
+}
diff --git a/android/TerminalApp/java/com/android/virtualization/terminal/DisplayProvider.kt b/android/TerminalApp/java/com/android/virtualization/terminal/DisplayProvider.kt
index fed8e5a..a04e056 100644
--- a/android/TerminalApp/java/com/android/virtualization/terminal/DisplayProvider.kt
+++ b/android/TerminalApp/java/com/android/virtualization/terminal/DisplayProvider.kt
@@ -39,7 +39,10 @@
     private val mainView: SurfaceView,
     private val cursorView: SurfaceView,
 ) {
-    private val virtService: IVirtualizationServiceInternal
+    private val virtService: IVirtualizationServiceInternal by lazy {
+        val b = ServiceManager.waitForService("android.system.virtualizationservice")
+        IVirtualizationServiceInternal.Stub.asInterface(b)
+    }
     private var cursorHandler: CursorHandler? = null
 
     init {
@@ -50,14 +53,6 @@
         cursorView.holder.setFormat(PixelFormat.RGBA_8888)
         // TODO: do we need this z-order?
         cursorView.setZOrderMediaOverlay(true)
-        val b = ServiceManager.waitForService("android.system.virtualizationservice")
-        virtService = IVirtualizationServiceInternal.Stub.asInterface(b)
-        try {
-            // To ensure that the previous display service is removed.
-            virtService.clearDisplayService()
-        } catch (e: RemoteException) {
-            throw RuntimeException("Failed to clear prior display service", e)
-        }
     }
 
     fun notifyDisplayIsGoingToInvisible() {
diff --git a/android/TerminalApp/java/com/android/virtualization/terminal/MainActivity.kt b/android/TerminalApp/java/com/android/virtualization/terminal/MainActivity.kt
index 1ae6ec5..71f10f9 100644
--- a/android/TerminalApp/java/com/android/virtualization/terminal/MainActivity.kt
+++ b/android/TerminalApp/java/com/android/virtualization/terminal/MainActivity.kt
@@ -56,6 +56,7 @@
 import androidx.activity.result.contract.ActivityResultContracts.StartActivityForResult
 import com.android.internal.annotations.VisibleForTesting
 import com.android.microdroid.test.common.DeviceProperties
+import com.android.system.virtualmachine.flags.Flags.terminalGuiSupport
 import com.android.virtualization.terminal.CertificateUtils.createOrGetKey
 import com.android.virtualization.terminal.CertificateUtils.writeCertificateToFile
 import com.android.virtualization.terminal.ErrorActivity.Companion.start
@@ -87,6 +88,7 @@
     private val bootCompleted = ConditionVariable()
     private lateinit var manageExternalStorageActivityResultLauncher: ActivityResultLauncher<Intent>
     private lateinit var modifierKeysController: ModifierKeysController
+    private var displayMenu: MenuItem? = null
 
     override fun onCreate(savedInstanceState: Bundle?) {
         super.onCreate(savedInstanceState)
@@ -256,6 +258,10 @@
                                     Trace.endAsyncSection("executeTerminal", 0)
                                     findViewById<View?>(R.id.boot_progress).visibility = View.GONE
                                     terminalContainer.visibility = View.VISIBLE
+                                    if (terminalGuiSupport()) {
+                                        displayMenu?.setVisible(true)
+                                        displayMenu?.setEnabled(true)
+                                    }
                                     bootCompleted.open()
                                     modifierKeysController.update()
                                     terminalView.mapTouchToMouseEvent()
@@ -344,6 +350,11 @@
 
     override fun onCreateOptionsMenu(menu: Menu?): Boolean {
         menuInflater.inflate(R.menu.main_menu, menu)
+        displayMenu =
+            menu?.findItem(R.id.menu_item_display).also {
+                it?.setVisible(terminalGuiSupport())
+                it?.setEnabled(false)
+            }
         return true
     }
 
@@ -353,6 +364,10 @@
             val intent = Intent(this, SettingsActivity::class.java)
             this.startActivity(intent)
             return true
+        } else if (id == R.id.menu_item_display) {
+            val intent = Intent(this, DisplayActivity::class.java)
+            this.startActivity(intent)
+            return true
         }
         return super.onOptionsItemSelected(item)
     }
diff --git a/android/TerminalApp/java/com/android/virtualization/terminal/VmLauncherService.kt b/android/TerminalApp/java/com/android/virtualization/terminal/VmLauncherService.kt
index f8b1b45..346056a 100644
--- a/android/TerminalApp/java/com/android/virtualization/terminal/VmLauncherService.kt
+++ b/android/TerminalApp/java/com/android/virtualization/terminal/VmLauncherService.kt
@@ -36,6 +36,7 @@
 import android.system.virtualmachine.VirtualMachineException
 import android.util.Log
 import android.widget.Toast
+import com.android.system.virtualmachine.flags.Flags.terminalGuiSupport
 import com.android.virtualization.terminal.MainActivity.Companion.TAG
 import com.android.virtualization.terminal.Runner.Companion.create
 import com.android.virtualization.terminal.VmLauncherService.VmLauncherServiceCallback
@@ -217,6 +218,21 @@
             changed = true
         }
 
+        // TODO(jeongik): let it configurable
+        if (terminalGuiSupport()) {
+            builder
+                .setDisplayConfig(
+                    VirtualMachineCustomImageConfig.DisplayConfig.Builder()
+                        .setWidth(1920)
+                        .setHeight(1080)
+                        .build()
+                )
+                .useKeyboard(true)
+                .useMouse(true)
+                .useTouch(true)
+            changed = true
+        }
+
         val image = InstalledImage.getDefault(this)
         if (image.hasBackup()) {
             val backup = image.backupFile
diff --git a/android/TerminalApp/proguard.flags b/android/TerminalApp/proguard.flags
index 88b8a9c..04a2140 100644
--- a/android/TerminalApp/proguard.flags
+++ b/android/TerminalApp/proguard.flags
@@ -4,7 +4,10 @@
 -keepattributes Signature
 
 # For using GSON @Expose annotation
--keepattributes *Annotation*
+-keepattributes RuntimeVisibleAnnotations,
+                RuntimeVisibleParameterAnnotations,
+                RuntimeVisibleTypeAnnotations,
+                AnnotationDefault
 
 # Gson specific classes
 -dontwarn sun.misc.**
diff --git a/android/TerminalApp/res/drawable/ic_display.xml b/android/TerminalApp/res/drawable/ic_display.xml
new file mode 100644
index 0000000..86bdb5d
--- /dev/null
+++ b/android/TerminalApp/res/drawable/ic_display.xml
@@ -0,0 +1,26 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--  Copyright 2025 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+ -->
+
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportWidth="960"
+    android:viewportHeight="960"
+    android:tint="?attr/colorControlNormal">
+    <path
+        android:fillColor="@android:color/white"
+        android:pathData="M240,840L240,760L280,720L160,720Q127,720 103.5,696.5Q80,673 80,640L80,200Q80,167 103.5,143.5Q127,120 160,120L800,120Q833,120 856.5,143.5Q880,167 880,200L880,640Q880,673 856.5,696.5Q833,720 800,720L680,720L720,760L720,840L240,840ZM160,640L800,640Q800,640 800,640Q800,640 800,640L800,200Q800,200 800,200Q800,200 800,200L160,200Q160,200 160,200Q160,200 160,200L160,640Q160,640 160,640Q160,640 160,640ZM160,640Q160,640 160,640Q160,640 160,640L160,200Q160,200 160,200Q160,200 160,200L160,200Q160,200 160,200Q160,200 160,200L160,640Q160,640 160,640Q160,640 160,640Z" />
+</vector>
\ No newline at end of file
diff --git a/android/TerminalApp/res/layout/activity_display.xml b/android/TerminalApp/res/layout/activity_display.xml
new file mode 100644
index 0000000..48e49fe
--- /dev/null
+++ b/android/TerminalApp/res/layout/activity_display.xml
@@ -0,0 +1,45 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--  Copyright 2024 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+ -->
+<merge xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:app="http://schemas.android.com/apk/res-auto"
+    xmlns:tools="http://schemas.android.com/tools"
+    android:layout_width="match_parent"
+    android:layout_height="match_parent"
+    tools:context=".DisplayActivity">
+    <View
+        android:id="@+id/background_touch_view"
+        android:layout_width="match_parent"
+        android:layout_height="match_parent"
+        />
+    <!-- the size should be match_parent -->
+    <SurfaceView
+        android:id="@+id/surface_view"
+        android:layout_width="1920px"
+        android:layout_height="1080px"
+        android:focusable="true"
+        android:focusableInTouchMode="true"
+        android:focusedByDefault="true"
+        android:defaultFocusHighlightEnabled="true">
+        <requestFocus />
+    </SurfaceView>
+    <!-- A cursor size in virtio-gpu spec is always 64x64 -->
+    <SurfaceView
+        android:id="@+id/cursor_surface_view"
+        android:layout_width="64px"
+        android:layout_height="64px">
+    </SurfaceView>
+
+</merge>
\ No newline at end of file
diff --git a/android/TerminalApp/res/menu/main_menu.xml b/android/TerminalApp/res/menu/main_menu.xml
index 42ba85d..dbb788c 100644
--- a/android/TerminalApp/res/menu/main_menu.xml
+++ b/android/TerminalApp/res/menu/main_menu.xml
@@ -20,4 +20,9 @@
         android:icon="@drawable/ic_settings"
         android:title="@string/action_settings"
         app:showAsAction="always"/>
+    <item android:id="@+id/menu_item_display"
+        android:icon="@drawable/ic_display"
+        android:enabled="false"
+        android:title="@string/action_display"
+        app:showAsAction="always"/>
 </menu>
diff --git a/android/TerminalApp/res/values/strings.xml b/android/TerminalApp/res/values/strings.xml
index bdebb83..44009c3 100644
--- a/android/TerminalApp/res/values/strings.xml
+++ b/android/TerminalApp/res/values/strings.xml
@@ -55,6 +55,9 @@
     <!-- Action bar icon name for the settings view CHAR LIMIT=none] -->
     <string name="action_settings">Settings</string>
 
+    <!-- Action bar icon name for showing the display activity CHAR LIMIT=none] -->
+    <string name="action_display">Display</string>
+
     <!-- Toast message to notify that preparing terminal to start [CHAR LIMIT=none] -->
     <string name="vm_creation_message">Preparing terminal</string>
     <!-- Toast message to notify that terminal is stopping [CHAR LIMIT=none] -->
diff --git a/android/TerminalApp/res/values/styles.xml b/android/TerminalApp/res/values/styles.xml
index 3fb8e7d..13f070f 100644
--- a/android/TerminalApp/res/values/styles.xml
+++ b/android/TerminalApp/res/values/styles.xml
@@ -27,4 +27,15 @@
     <style name="VmTerminalAppTheme" parent="@style/Theme.Material3.DayNight.NoActionBar">
         <item name="android:windowLightStatusBar" tools:targetApi="m">?android:attr/isLightTheme</item>
     </style>
+    <style name="FullscreenTheme" parent="@style/Theme.Material3.DayNight.NoActionBar">
+        <item name="android:navigationBarColor">
+            @android:color/transparent
+        </item>
+        <item name="android:statusBarColor">
+            @android:color/transparent
+        </item>
+        <item name="android:windowLayoutInDisplayCutoutMode">
+            always
+        </item>
+    </style>
 </resources>
diff --git a/android/vm/src/run.rs b/android/vm/src/run.rs
index 0037327..0855fa0 100644
--- a/android/vm/src/run.rs
+++ b/android/vm/src/run.rs
@@ -339,11 +339,10 @@
     } else {
         None
     };
+    let vm = VmInstance::create(service, config, console_out, console_in, log, dump_dt)
+        .context("Failed to create VM")?;
     let callback = Box::new(Callback {});
-    let vm =
-        VmInstance::create(service, config, console_out, console_in, log, dump_dt, Some(callback))
-            .context("Failed to create VM")?;
-    vm.start().context("Failed to start VM")?;
+    vm.start(Some(callback)).context("Failed to start VM")?;
 
     let debug_level = get_debug_level(config).unwrap_or(DebugLevel::NONE);
 
diff --git a/build/debian/fai_config/package_config/AVF b/build/debian/fai_config/package_config/AVF
index 98b558b..bf51fdb 100644
--- a/build/debian/fai_config/package_config/AVF
+++ b/build/debian/fai_config/package_config/AVF
@@ -8,3 +8,5 @@
 forwarder-guest
 forwarder-guest-launcher
 shutdown-runner
+weston
+xwayland
diff --git a/guest/microdroid_launcher/Android.bp b/guest/microdroid_launcher/Android.bp
index 42c18cb..893ee98 100644
--- a/guest/microdroid_launcher/Android.bp
+++ b/guest/microdroid_launcher/Android.bp
@@ -12,5 +12,10 @@
         "libdl_android",
         "liblog",
     ],
+    static_libs: [
+        "libapexutil",
+        "libprotobuf-cpp-lite",
+        "lib_apex_manifest_proto_lite",
+    ],
     header_libs: ["vm_payload_headers"],
 }
diff --git a/guest/microdroid_launcher/main.cpp b/guest/microdroid_launcher/main.cpp
index 5ae9956..8e3d4e4 100644
--- a/guest/microdroid_launcher/main.cpp
+++ b/guest/microdroid_launcher/main.cpp
@@ -16,7 +16,9 @@
 
 #include <android-base/logging.h>
 #include <android-base/result.h>
+#include <android-base/strings.h>
 #include <android/dlext.h>
+#include <apexutil.h>
 #include <dlfcn.h>
 
 #include <cstdlib>
@@ -25,8 +27,11 @@
 
 #include "vm_main.h"
 
+using android::apex::GetActivePackages;
 using android::base::Error;
+using android::base::Join;
 using android::base::Result;
+using android::base::StringReplace;
 
 extern "C" {
 enum {
@@ -43,6 +48,8 @@
 extern bool android_link_namespaces(struct android_namespace_t* from,
                                     struct android_namespace_t* to,
                                     const char* shared_libs_sonames);
+
+extern struct android_namespace_t* android_get_exported_namespace(const char* name);
 } // extern "C"
 
 static Result<void*> load(const std::string& libname);
@@ -108,6 +115,32 @@
         return Error() << "Failed to link namespace: " << dlerror();
     }
 
+    // Make libraries provided by APEXes available to the payload
+    for (const auto& [_path, manifest] : GetActivePackages("/apex")) {
+        std::string namespace_name = StringReplace(manifest.name(), ".", "_", /* all */ true);
+        android_namespace_t* apex_ns = android_get_exported_namespace(namespace_name.c_str());
+        if (apex_ns == nullptr) {
+            // This means the namespace is configured as 'visible=false'. We can't link to an
+            // invisible namespace.
+            continue;
+        }
+
+        std::vector<std::string> libs = {manifest.providenativelibs().begin(),
+                                         manifest.providenativelibs().end()};
+        if (libs.size() == 0) {
+            continue;
+        }
+        std::string shared_lib_sonames = Join(libs, ':');
+
+        if (!android_link_namespaces(new_ns, apex_ns, shared_lib_sonames.c_str())) {
+            return Error() << "Failed to link APEX namespace " << namespace_name << ": "
+                           << dlerror();
+        }
+
+        LOG(INFO) << "Linked APEX namespace " << namespace_name << " with shared libs "
+                  << shared_lib_sonames;
+    }
+
     const android_dlextinfo info = {
             .flags = ANDROID_DLEXT_USE_NAMESPACE,
             .library_namespace = new_ns,
diff --git a/guest/microdroid_manager/src/vm_secret.rs b/guest/microdroid_manager/src/vm_secret.rs
index 56b3482..5999122 100644
--- a/guest/microdroid_manager/src/vm_secret.rs
+++ b/guest/microdroid_manager/src/vm_secret.rs
@@ -171,7 +171,11 @@
             return Err(anyhow!("Rollback protected data is not available with V1 secrets"));
         };
         let payload_id = sha::sha512(instance_id);
-        secretkeeper_session.get_secret(payload_id)
+        secretkeeper_session.get_secret(payload_id).or_else(|e| {
+            log::info!("Secretkeeper get failed with {e:?}. Refreshing connection & retrying!");
+            secretkeeper_session.refresh()?;
+            secretkeeper_session.get_secret(payload_id)
+        })
     }
 
     pub fn write_payload_data_rp(&self, data: &[u8; SECRET_SIZE]) -> Result<()> {
@@ -180,7 +184,12 @@
             return Err(anyhow!("Rollback protected data is not available with V1 secrets"));
         };
         let payload_id = sha::sha512(instance_id);
-        secretkeeper_session.store_secret(payload_id, data)
+        if let Err(e) = secretkeeper_session.store_secret(payload_id, data.clone()) {
+            log::info!("Secretkeeper store failed with {e:?}. Refreshing connection & retrying!");
+            secretkeeper_session.refresh()?;
+            secretkeeper_session.store_secret(payload_id, data)?;
+        }
+        Ok(())
     }
 }
 
@@ -276,6 +285,11 @@
         Ok(Self { session, sealing_policy })
     }
 
+    fn refresh(&self) -> Result<()> {
+        let mut session = self.session.lock().unwrap();
+        Ok(session.refresh()?)
+    }
+
     fn store_secret(&self, id: [u8; ID_SIZE], secret: Zeroizing<[u8; SECRET_SIZE]>) -> Result<()> {
         let store_request = StoreSecretRequest {
             id: Id(id),
diff --git a/guest/pvmfw/src/dice.rs b/guest/pvmfw/src/dice.rs
index 78bd6b8..4df10b3 100644
--- a/guest/pvmfw/src/dice.rs
+++ b/guest/pvmfw/src/dice.rs
@@ -156,9 +156,7 @@
     fn generate_config_descriptor(&self, instance_hash: Option<Hash>) -> Result<Vec<u8>> {
         let mut config = Vec::with_capacity(4);
         config.push((cbor!(COMPONENT_NAME_KEY)?, cbor!("vm_entry")?));
-        if cfg!(dice_changes) {
-            config.push((cbor!(SECURITY_VERSION_KEY)?, cbor!(self.security_version)?));
-        }
+        config.push((cbor!(SECURITY_VERSION_KEY)?, cbor!(self.security_version)?));
         if self.rkp_vm_marker {
             config.push((cbor!(RKP_VM_MARKER_KEY)?, Value::Null))
         }
@@ -245,14 +243,7 @@
         assert_eq!(config_map.get(&COMPONENT_NAME_KEY).unwrap().as_text().unwrap(), "vm_entry");
         assert_eq!(config_map.get(&COMPONENT_VERSION_KEY), None);
         assert_eq!(config_map.get(&RESETTABLE_KEY), None);
-        if cfg!(dice_changes) {
-            assert_eq!(
-                config_map.get(&SECURITY_VERSION_KEY).unwrap().as_integer().unwrap(),
-                42.into()
-            );
-        } else {
-            assert_eq!(config_map.get(&SECURITY_VERSION_KEY), None);
-        }
+        assert_eq!(config_map.get(&SECURITY_VERSION_KEY).unwrap().as_integer().unwrap(), 42.into());
         assert_eq!(config_map.get(&RKP_VM_MARKER_KEY), None);
     }
 
diff --git a/guest/pvmfw/src/fdt.rs b/guest/pvmfw/src/fdt.rs
index 29212f9..818d342 100644
--- a/guest/pvmfw/src/fdt.rs
+++ b/guest/pvmfw/src/fdt.rs
@@ -112,6 +112,24 @@
     Ok(None)
 }
 
+/// Read /avf/untrusted/instance-id, if present.
+pub fn read_instance_id(fdt: &Fdt) -> libfdt::Result<Option<&[u8]>> {
+    read_avf_untrusted_prop(fdt, c"instance-id")
+}
+
+/// Read /avf/untrusted/defer-rollback-protection, if present.
+pub fn read_defer_rollback_protection(fdt: &Fdt) -> libfdt::Result<Option<&[u8]>> {
+    read_avf_untrusted_prop(fdt, c"defer-rollback-protection")
+}
+
+fn read_avf_untrusted_prop<'a>(fdt: &'a Fdt, prop: &CStr) -> libfdt::Result<Option<&'a [u8]>> {
+    if let Some(node) = fdt.node(c"/avf/untrusted")? {
+        node.getprop(prop)
+    } else {
+        Ok(None)
+    }
+}
+
 fn patch_initrd_range(fdt: &mut Fdt, initrd_range: &Range<usize>) -> libfdt::Result<()> {
     let start = u32::try_from(initrd_range.start).unwrap();
     let end = u32::try_from(initrd_range.end).unwrap();
diff --git a/guest/pvmfw/src/main.rs b/guest/pvmfw/src/main.rs
index 0a3dca6..afa64e0 100644
--- a/guest/pvmfw/src/main.rs
+++ b/guest/pvmfw/src/main.rs
@@ -35,22 +35,20 @@
 use crate::bcc::Bcc;
 use crate::dice::PartialInputs;
 use crate::entry::RebootReason;
-use crate::fdt::{modify_for_next_stage, sanitize_device_tree};
+use crate::fdt::{modify_for_next_stage, read_instance_id, sanitize_device_tree};
 use crate::rollback::perform_rollback_protection;
 use alloc::borrow::Cow;
 use alloc::boxed::Box;
 use bssl_avf::Digester;
 use diced_open_dice::{bcc_handover_parse, DiceArtifacts, DiceContext, Hidden, VM_KEY_ALGORITHM};
-use libfdt::{Fdt, FdtNode};
+use libfdt::Fdt;
 use log::{debug, error, info, trace, warn};
 use pvmfw_avb::verify_payload;
 use pvmfw_avb::DebugLevel;
 use pvmfw_embedded_key::PUBLIC_KEY;
-use vmbase::fdt::pci::{PciError, PciInfo};
 use vmbase::heap;
-use vmbase::memory::{flush, init_shared_pool, SIZE_4KB};
+use vmbase::memory::{flush, SIZE_4KB};
 use vmbase::rand;
-use vmbase::virtio::pci;
 
 fn main<'a>(
     untrusted_fdt: &mut Fdt,
@@ -77,8 +75,6 @@
     })?;
     trace!("BCC: {bcc_handover:x?}");
 
-    let cdi_seal = bcc_handover.cdi_seal();
-
     let bcc = Bcc::new(bcc_handover.bcc()).map_err(|e| {
         error!("{e}");
         RebootReason::InvalidBcc
@@ -102,19 +98,8 @@
     }
 
     let guest_page_size = verified_boot_data.page_size.unwrap_or(SIZE_4KB);
-    let fdt_info = sanitize_device_tree(untrusted_fdt, vm_dtbo, vm_ref_dt, guest_page_size)?;
+    let _ = sanitize_device_tree(untrusted_fdt, vm_dtbo, vm_ref_dt, guest_page_size)?;
     let fdt = untrusted_fdt; // DT has now been sanitized.
-    let pci_info = PciInfo::from_fdt(fdt).map_err(handle_pci_error)?;
-    debug!("PCI: {:#x?}", pci_info);
-    // Set up PCI bus for VirtIO devices.
-    let mut pci_root = pci::initialize(pci_info).map_err(|e| {
-        error!("Failed to initialize PCI: {e}");
-        RebootReason::InternalError
-    })?;
-    init_shared_pool(fdt_info.swiotlb_info.fixed_range()).map_err(|e| {
-        error!("Failed to initialize shared pool: {e}");
-        RebootReason::InternalError
-    })?;
 
     let next_bcc_size = guest_page_size;
     let next_bcc = heap::aligned_boxed_slice(next_bcc_size, guest_page_size).ok_or_else(|| {
@@ -129,13 +114,12 @@
         RebootReason::InternalError
     })?;
 
-    let instance_hash = Some(salt_from_instance_id(fdt)?);
+    let instance_hash = salt_from_instance_id(fdt)?;
     let (new_instance, salt, defer_rollback_protection) = perform_rollback_protection(
         fdt,
         &verified_boot_data,
         &dice_inputs,
-        &mut pci_root,
-        cdi_seal,
+        bcc_handover.cdi_seal(),
         instance_hash,
     )?;
     trace!("Got salt for instance: {salt:x?}");
@@ -204,8 +188,14 @@
 
 // Get the "salt" which is one of the input for DICE derivation.
 // This provides differentiation of secrets for different VM instances with same payloads.
-fn salt_from_instance_id(fdt: &Fdt) -> Result<Hidden, RebootReason> {
-    let id = instance_id(fdt)?;
+fn salt_from_instance_id(fdt: &Fdt) -> Result<Option<Hidden>, RebootReason> {
+    let Some(id) = read_instance_id(fdt).map_err(|e| {
+        error!("Failed to get instance-id in DT: {e}");
+        RebootReason::InvalidFdt
+    })?
+    else {
+        return Ok(None);
+    };
     let salt = Digester::sha512()
         .digest(&[&b"InstanceId:"[..], id].concat())
         .map_err(|e| {
@@ -214,46 +204,5 @@
         })?
         .try_into()
         .map_err(|_| RebootReason::InternalError)?;
-    Ok(salt)
-}
-
-fn instance_id(fdt: &Fdt) -> Result<&[u8], RebootReason> {
-    let node = avf_untrusted_node(fdt)?;
-    let id = node.getprop(c"instance-id").map_err(|e| {
-        error!("Failed to get instance-id in DT: {e}");
-        RebootReason::InvalidFdt
-    })?;
-    id.ok_or_else(|| {
-        error!("Missing instance-id");
-        RebootReason::InvalidFdt
-    })
-}
-
-fn avf_untrusted_node(fdt: &Fdt) -> Result<FdtNode, RebootReason> {
-    let node = fdt.node(c"/avf/untrusted").map_err(|e| {
-        error!("Failed to get /avf/untrusted node: {e}");
-        RebootReason::InvalidFdt
-    })?;
-    node.ok_or_else(|| {
-        error!("/avf/untrusted node is missing in DT");
-        RebootReason::InvalidFdt
-    })
-}
-
-/// Logs the given PCI error and returns the appropriate `RebootReason`.
-fn handle_pci_error(e: PciError) -> RebootReason {
-    error!("{}", e);
-    match e {
-        PciError::FdtErrorPci(_)
-        | PciError::FdtNoPci
-        | PciError::FdtErrorReg(_)
-        | PciError::FdtMissingReg
-        | PciError::FdtRegEmpty
-        | PciError::FdtRegMissingSize
-        | PciError::CamWrongSize(_)
-        | PciError::FdtErrorRanges(_)
-        | PciError::FdtMissingRanges
-        | PciError::RangeAddressMismatch { .. }
-        | PciError::NoSuitableRange => RebootReason::InvalidFdt,
-    }
+    Ok(Some(salt))
 }
diff --git a/guest/pvmfw/src/rollback.rs b/guest/pvmfw/src/rollback.rs
index f7723d7..74b2cd8 100644
--- a/guest/pvmfw/src/rollback.rs
+++ b/guest/pvmfw/src/rollback.rs
@@ -16,16 +16,20 @@
 
 use crate::dice::PartialInputs;
 use crate::entry::RebootReason;
+use crate::fdt::read_defer_rollback_protection;
 use crate::instance::EntryBody;
 use crate::instance::Error as InstanceError;
 use crate::instance::{get_recorded_entry, record_instance_entry};
 use diced_open_dice::Hidden;
-use libfdt::{Fdt, FdtNode};
+use libfdt::Fdt;
 use log::{error, info};
 use pvmfw_avb::Capability;
 use pvmfw_avb::VerifiedBootData;
 use virtio_drivers::transport::pci::bus::PciRoot;
+use vmbase::fdt::{pci::PciInfo, SwiotlbInfo};
+use vmbase::memory::init_shared_pool;
 use vmbase::rand;
+use vmbase::virtio::pci;
 
 /// Performs RBP based on the input payload, current DICE chain, and host-controlled platform.
 ///
@@ -37,7 +41,6 @@
     fdt: &Fdt,
     verified_boot_data: &VerifiedBootData,
     dice_inputs: &PartialInputs,
-    pci_root: &mut PciRoot,
     cdi_seal: &[u8],
     instance_hash: Option<Hidden>,
 ) -> Result<(bool, Hidden, bool), RebootReason> {
@@ -53,7 +56,7 @@
         skip_rollback_protection()?;
         Ok((false, instance_hash.unwrap(), false))
     } else {
-        perform_legacy_rollback_protection(dice_inputs, pci_root, cdi_seal, instance_hash)
+        perform_legacy_rollback_protection(fdt, dice_inputs, cdi_seal, instance_hash)
     }
 }
 
@@ -92,17 +95,18 @@
 
 /// Performs RBP using instance.img where updates require clearing old entries, causing new CDIs.
 fn perform_legacy_rollback_protection(
+    fdt: &Fdt,
     dice_inputs: &PartialInputs,
-    pci_root: &mut PciRoot,
     cdi_seal: &[u8],
     instance_hash: Option<Hidden>,
 ) -> Result<(bool, Hidden, bool), RebootReason> {
     info!("Fallback to instance.img based rollback checks");
-    let (recorded_entry, mut instance_img, header_index) = get_recorded_entry(pci_root, cdi_seal)
-        .map_err(|e| {
-        error!("Failed to get entry from instance.img: {e}");
-        RebootReason::InternalError
-    })?;
+    let mut pci_root = initialize_instance_img_device(fdt)?;
+    let (recorded_entry, mut instance_img, header_index) =
+        get_recorded_entry(&mut pci_root, cdi_seal).map_err(|e| {
+            error!("Failed to get entry from instance.img: {e}");
+            RebootReason::InternalError
+        })?;
     let (new_instance, salt) = if let Some(entry) = recorded_entry {
         check_dice_measurements_match_entry(dice_inputs, &entry)?;
         let salt = instance_hash.unwrap_or(entry.salt);
@@ -155,24 +159,34 @@
 }
 
 fn should_defer_rollback_protection(fdt: &Fdt) -> Result<bool, RebootReason> {
-    let node = avf_untrusted_node(fdt)?;
-    let defer_rbp = node
-        .getprop(c"defer-rollback-protection")
-        .map_err(|e| {
-            error!("Failed to get defer-rollback-protection property in DT: {e}");
-            RebootReason::InvalidFdt
-        })?
-        .is_some();
-    Ok(defer_rbp)
-}
-
-fn avf_untrusted_node(fdt: &Fdt) -> Result<FdtNode, RebootReason> {
-    let node = fdt.node(c"/avf/untrusted").map_err(|e| {
-        error!("Failed to get /avf/untrusted node: {e}");
+    let defer_rbp = read_defer_rollback_protection(fdt).map_err(|e| {
+        error!("Failed to get defer-rollback-protection property in DT: {e}");
         RebootReason::InvalidFdt
     })?;
-    node.ok_or_else(|| {
-        error!("/avf/untrusted node is missing in DT");
+    Ok(defer_rbp.is_some())
+}
+
+/// Set up PCI bus and VirtIO-blk device containing the instance.img partition.
+fn initialize_instance_img_device(fdt: &Fdt) -> Result<PciRoot, RebootReason> {
+    let pci_info = PciInfo::from_fdt(fdt).map_err(|e| {
+        error!("Failed to detect PCI from DT: {e}");
         RebootReason::InvalidFdt
-    })
+    })?;
+    let swiotlb_range = SwiotlbInfo::new_from_fdt(fdt)
+        .map_err(|e| {
+            error!("Failed to detect swiotlb from DT: {e}");
+            RebootReason::InvalidFdt
+        })?
+        .and_then(|info| info.fixed_range());
+
+    let pci_root = pci::initialize(pci_info).map_err(|e| {
+        error!("Failed to initialize PCI: {e}");
+        RebootReason::InternalError
+    })?;
+    init_shared_pool(swiotlb_range).map_err(|e| {
+        error!("Failed to initialize shared pool: {e}");
+        RebootReason::InternalError
+    })?;
+
+    Ok(pci_root)
 }
diff --git a/guest/rialto/tests/test.rs b/guest/rialto/tests/test.rs
index c94a0e3..c650046 100644
--- a/guest/rialto/tests/test.rs
+++ b/guest/rialto/tests/test.rs
@@ -54,13 +54,9 @@
 const INSTANCE_IMG_PATH: &str = "/data/local/tmp/rialto_test/arm64/instance.img";
 const TEST_CERT_CHAIN_PATH: &str = "testdata/rkp_cert_chain.der";
 
-#[cfg(dice_changes)]
 #[test]
 fn process_requests_in_protected_vm() -> Result<()> {
     if hypervisor_props::is_protected_vm_supported()? {
-        // The test is skipped if the feature flag |dice_changes| is not enabled, because when
-        // the flag is off, the DICE chain is truncated in the pvmfw, and the service VM cannot
-        // verify the chain due to the missing entries in the chain.
         check_processing_requests(VmType::ProtectedVm, None)
     } else {
         warn!("pVMs are not supported on device, skipping test");
@@ -342,7 +338,6 @@
         /* consoleIn */ None,
         log,
         /* dump_dt */ None,
-        None,
     )
     .context("Failed to create VM")
 }
diff --git a/guest/trusty/security_vm/launcher/src/main.rs b/guest/trusty/security_vm/launcher/src/main.rs
index 8dd7c43..2fedc16 100644
--- a/guest/trusty/security_vm/launcher/src/main.rs
+++ b/guest/trusty/security_vm/launcher/src/main.rs
@@ -101,10 +101,9 @@
         None, // console_out
         None, // log
         None, // dump_dt
-        None, // callback
     )
     .context("Failed to create VM")?;
-    vm.start().context("Failed to start VM")?;
+    vm.start(None /* callback */).context("Failed to start VM")?;
 
     println!("started trusty_security_vm_launcher VM");
     let death_reason = vm.wait_for_death();
diff --git a/libs/libavf/include/android/virtualization.h b/libs/libavf/include/android/virtualization.h
index ef57325..03d04a9 100644
--- a/libs/libavf/include/android/virtualization.h
+++ b/libs/libavf/include/android/virtualization.h
@@ -337,17 +337,30 @@
 int AVirtualMachine_createRaw(const AVirtualizationService* _Nonnull service,
                               AVirtualMachineRawConfig* _Nonnull config, int consoleOutFd,
                               int consoleInFd, int logFd,
-                              AVirtualMachine* _Null_unspecified* _Nonnull vm) __INTRODUCED_IN(36);
+                              const AVirtualMachine* _Null_unspecified* _Nonnull vm)
+        __INTRODUCED_IN(36);
+
+/**
+ * A callback to be called when virtual machine stops.
+ *
+ * \param vm stopped vm
+ * \param reason stop reason
+ */
+typedef void (*AVirtualMachine_stopCallback)(const AVirtualMachine* _Nonnull vm,
+                                             enum AVirtualMachineStopReason reason);
 
 /**
  * Start a virtual machine. `AVirtualMachine_start` is synchronous and blocks until the virtual
  * machine is initialized and free to start executing code, or until an error happens.
  *
  * \param vm a handle on a virtual machine.
+ * \param callback an optional callback to be called when VM stops.
  *
  * \return If successful, it returns 0. Otherwise, it returns `-EIO`.
  */
-int AVirtualMachine_start(AVirtualMachine* _Nonnull vm) __INTRODUCED_IN(36);
+int AVirtualMachine_start(const AVirtualMachine* _Nonnull vm,
+                          const AVirtualMachine_stopCallback _Nullable callback)
+        __INTRODUCED_IN(36);
 
 /**
  * Stop a virtual machine. Stopping a virtual machine is like pulling the plug on a real computer;
@@ -366,7 +379,7 @@
  *
  * \return If successful, it returns 0. Otherwise, it returns `-EIO`.
  */
-int AVirtualMachine_stop(AVirtualMachine* _Nonnull vm) __INTRODUCED_IN(36);
+int AVirtualMachine_stop(const AVirtualMachine* _Nonnull vm) __INTRODUCED_IN(36);
 
 /**
  * Open a vsock connection to the VM on the given port. The caller takes ownership of the returned
@@ -379,7 +392,8 @@
  *
  * \return If successful, it returns a valid file descriptor. Otherwise, it returns `-EIO`.
  */
-int AVirtualMachine_connectVsock(AVirtualMachine* _Nonnull vm, uint32_t port) __INTRODUCED_IN(36);
+int AVirtualMachine_connectVsock(const AVirtualMachine* _Nonnull vm, uint32_t port)
+        __INTRODUCED_IN(36);
 
 /**
  * Wait until a virtual machine stops or the given timeout elapses.
@@ -393,7 +407,7 @@
  *     sets `reason` and returns true.
  *   - If the timeout expired, it returns `false`.
  */
-bool AVirtualMachine_waitForStop(AVirtualMachine* _Nonnull vm,
+bool AVirtualMachine_waitForStop(const AVirtualMachine* _Nonnull vm,
                                  const struct timespec* _Nullable timeout,
                                  enum AVirtualMachineStopReason* _Nonnull reason)
         __INTRODUCED_IN(36);
@@ -408,6 +422,6 @@
  *
  * \param vm a handle on a virtual machine.
  */
-void AVirtualMachine_destroy(AVirtualMachine* _Nullable vm) __INTRODUCED_IN(36);
+void AVirtualMachine_destroy(const AVirtualMachine* _Nullable vm) __INTRODUCED_IN(36);
 
 __END_DECLS
diff --git a/libs/libavf/src/lib.rs b/libs/libavf/src/lib.rs
index 6532ace..256803f 100644
--- a/libs/libavf/src/lib.rs
+++ b/libs/libavf/src/lib.rs
@@ -19,6 +19,7 @@
 use std::os::fd::{FromRawFd, IntoRawFd};
 use std::os::raw::{c_char, c_int};
 use std::ptr;
+use std::sync::Arc;
 use std::time::Duration;
 
 use android_system_virtualizationservice::{
@@ -29,10 +30,10 @@
     },
     binder::{ParcelFileDescriptor, Strong},
 };
-use avf_bindgen::AVirtualMachineStopReason;
+use avf_bindgen::{AVirtualMachineStopReason, AVirtualMachine_stopCallback};
 use libc::timespec;
 use log::error;
-use vmclient::{DeathReason, VirtualizationService, VmInstance};
+use vmclient::{DeathReason, ErrorCode, VirtualizationService, VmCallback, VmInstance};
 
 /// Create a new virtual machine config object with no properties.
 #[no_mangle]
@@ -342,6 +343,49 @@
     }
 }
 
+struct LocalVmInstance {
+    vm: Arc<VmInstance>,
+    callback: AVirtualMachine_stopCallback,
+}
+
+impl VmCallback for LocalVmInstance {
+    fn on_payload_started(&self, _cid: i32) {
+        // Microdroid only. no-op.
+    }
+
+    fn on_payload_ready(&self, _cid: i32) {
+        // Microdroid only. no-op.
+    }
+
+    fn on_payload_finished(&self, _cid: i32, _exit_code: i32) {
+        // Microdroid only. no-op.
+    }
+
+    fn on_error(&self, _cid: i32, _error_code: ErrorCode, _message: &str) {
+        // Microdroid only. no-op.
+    }
+
+    fn on_died(&self, _cid: i32, death_reason: DeathReason) {
+        let Some(callback) = self.callback else {
+            return;
+        };
+        let stop_reason = death_reason_to_stop_reason(death_reason);
+        let vm_ptr: *const VmInstance = Arc::into_raw(Arc::clone(&self.vm));
+
+        // SAFETY: `callback` is assumed to be a valid, non-null function pointer passed by
+        // `AVirtualMachine_start`.
+        unsafe {
+            callback(vm_ptr.cast(), stop_reason);
+        }
+
+        // drop ptr after use.
+        // SAFETY: `vm_ptr` is a valid, non-null pointer casted above.
+        unsafe {
+            let _ = Arc::from_raw(vm_ptr);
+        }
+    }
+}
+
 /// Create a virtual machine with given `config`.
 ///
 /// # Safety
@@ -357,7 +401,7 @@
     console_out_fd: c_int,
     console_in_fd: c_int,
     log_fd: c_int,
-    vm_ptr: *mut *mut VmInstance,
+    vm_ptr: *mut *const VmInstance,
 ) -> c_int {
     // SAFETY: `service` is assumed to be a valid, non-null pointer returned by
     // `AVirtualizationService_create` or `AVirtualizationService_create_early`. It's the only
@@ -373,12 +417,11 @@
     let console_in = get_file_from_fd(console_in_fd);
     let log = get_file_from_fd(log_fd);
 
-    match VmInstance::create(service.as_ref(), &config, console_out, console_in, log, None, None) {
+    match VmInstance::create(service.as_ref(), &config, console_out, console_in, log, None) {
         Ok(vm) => {
             // SAFETY: `vm_ptr` is assumed to be a valid, non-null pointer to a mutable raw pointer.
-            // `vm` is the only reference here and `vm_ptr` takes ownership.
             unsafe {
-                *vm_ptr = Box::into_raw(Box::new(vm));
+                *vm_ptr = Arc::into_raw(Arc::new(vm));
             }
             0
         }
@@ -394,11 +437,20 @@
 /// # Safety
 /// `vm` must be a pointer returned by `AVirtualMachine_createRaw`.
 #[no_mangle]
-pub unsafe extern "C" fn AVirtualMachine_start(vm: *const VmInstance) -> c_int {
+pub unsafe extern "C" fn AVirtualMachine_start(
+    vm: *const VmInstance,
+    callback: AVirtualMachine_stopCallback,
+) -> c_int {
     // SAFETY: `vm` is assumed to be a valid, non-null pointer returned by
     // `AVirtualMachine_createRaw`. It's the only reference to the object.
-    let vm = unsafe { &*vm };
-    match vm.start() {
+    let vm = unsafe { Arc::from_raw(vm) };
+    let callback = callback.map(|_| {
+        let cb: Box<dyn VmCallback + Send + Sync> =
+            Box::new(LocalVmInstance { vm: Arc::clone(&vm), callback });
+        cb
+    });
+
+    match vm.start(callback) {
         Ok(_) => 0,
         Err(e) => {
             error!("AVirtualMachine_start failed: {e:?}");
@@ -509,12 +561,12 @@
 /// `vm` must be a pointer returned by `AVirtualMachine_createRaw`. `vm` must not be reused after
 /// deletion.
 #[no_mangle]
-pub unsafe extern "C" fn AVirtualMachine_destroy(vm: *mut VmInstance) {
+pub unsafe extern "C" fn AVirtualMachine_destroy(vm: *const VmInstance) {
     if !vm.is_null() {
         // SAFETY: `vm` is assumed to be a valid, non-null pointer returned by
         // AVirtualMachine_create. It's the only reference to the object.
         unsafe {
-            let _ = Box::from_raw(vm);
+            let _ = Arc::from_raw(vm);
         }
     }
 }
diff --git a/libs/libcompos_common/compos_client.rs b/libs/libcompos_common/compos_client.rs
index 6872582..c2b4936 100644
--- a/libs/libcompos_common/compos_client.rs
+++ b/libs/libcompos_common/compos_client.rs
@@ -148,19 +148,14 @@
 
         // Let logs go to logcat.
         let (console_fd, log_fd) = (None, None);
-        let callback = Box::new(Callback {});
         let instance = VmInstance::create(
-            service,
-            &config,
-            console_fd,
-            /* console_in_fd */ None,
-            log_fd,
+            service, &config, console_fd, /* console_in_fd */ None, log_fd,
             /* dump_dt */ None,
-            Some(callback),
         )
         .context("Failed to create VM")?;
 
-        instance.start()?;
+        let callback = Box::new(Callback {});
+        instance.start(Some(callback))?;
 
         let ready = instance.wait_until_ready(TIMEOUTS.vm_max_time_to_ready);
         if ready == Err(VmWaitError::Finished) && debug_level != DebugLevel::NONE {
diff --git a/libs/libservice_vm_manager/src/lib.rs b/libs/libservice_vm_manager/src/lib.rs
index 0f322bb..77e7a4a 100644
--- a/libs/libservice_vm_manager/src/lib.rs
+++ b/libs/libservice_vm_manager/src/lib.rs
@@ -152,7 +152,7 @@
         let vsock_listener = VsockListener::bind_with_cid_port(VMADDR_CID_HOST, vm_type.port())?;
 
         // Starts the service VM.
-        vm.start().context("Failed to start service VM")?;
+        vm.start(None).context("Failed to start service VM")?;
         info!("Service VM started");
 
         // Accepts the connection from the service VM.
@@ -245,8 +245,7 @@
     let console_in = None;
     let log = Some(android_log_fd()?);
     let dump_dt = None;
-    let callback = None;
-    VmInstance::create(service.as_ref(), &config, console_out, console_in, log, dump_dt, callback)
+    VmInstance::create(service.as_ref(), &config, console_out, console_in, log, dump_dt)
         .context("Failed to create service VM")
 }
 
diff --git a/libs/libvmclient/src/lib.rs b/libs/libvmclient/src/lib.rs
index 8dd3cd3..2c6abb5 100644
--- a/libs/libvmclient/src/lib.rs
+++ b/libs/libvmclient/src/lib.rs
@@ -209,7 +209,6 @@
         console_in: Option<File>,
         log: Option<File>,
         dump_dt: Option<File>,
-        callback: Option<Box<dyn VmCallback + Send + Sync>>,
     ) -> BinderResult<Self> {
         let console_out = console_out.map(ParcelFileDescriptor::new);
         let console_in = console_in.map(ParcelFileDescriptor::new);
@@ -226,20 +225,19 @@
 
         let cid = vm.getCid()?;
 
-        // Register callback before starting VM, in case it dies immediately.
         let state = Arc::new(Monitor::new(VmState::default()));
-        let callback = BnVirtualMachineCallback::new_binder(
-            VirtualMachineCallback { state: state.clone(), client_callback: callback },
-            BinderFeatures::default(),
-        );
-        vm.registerCallback(&callback)?;
         let death_recipient = wait_for_binder_death(&mut vm.as_binder(), state.clone())?;
 
         Ok(Self { vm, cid, state, _death_recipient: death_recipient })
     }
 
     /// Starts the VM.
-    pub fn start(&self) -> BinderResult<()> {
+    pub fn start(&self, callback: Option<Box<dyn VmCallback + Send + Sync>>) -> BinderResult<()> {
+        let callback = BnVirtualMachineCallback::new_binder(
+            VirtualMachineCallback { state: self.state.clone(), client_callback: callback },
+            BinderFeatures::default(),
+        );
+        self.vm.registerCallback(&callback)?;
         self.vm.start()
     }
 
diff --git a/microfuchsia/microfuchsiad/src/instance_starter.rs b/microfuchsia/microfuchsiad/src/instance_starter.rs
index e3c4e8d..7c4f32d 100644
--- a/microfuchsia/microfuchsiad/src/instance_starter.rs
+++ b/microfuchsia/microfuchsiad/src/instance_starter.rs
@@ -96,7 +96,6 @@
             console_in,
             /* log= */ None,
             /* dump_dt= */ None,
-            None,
         )
         .context("Failed to create VM")?;
         if let Some(pty) = &pty {
@@ -105,7 +104,7 @@
                 .setHostConsoleName(&pty.follower_name)
                 .context("Setting host console name")?;
         }
-        vm_instance.start().context("Starting VM")?;
+        vm_instance.start(None).context("Starting VM")?;
 
         Ok(MicrofuchsiaInstance {
             _vm_instance: vm_instance,
diff --git a/tests/backcompat_test/src/main.rs b/tests/backcompat_test/src/main.rs
index 4113881..a7fb074 100644
--- a/tests/backcompat_test/src/main.rs
+++ b/tests/backcompat_test/src/main.rs
@@ -117,10 +117,9 @@
         /* consoleIn */ None,
         None,
         Some(dump_dt),
-        None,
     )
     .context("Failed to create VM")?;
-    vm.start().context("Failed to start VM")?;
+    vm.start(None).context("Failed to start VM")?;
     info!("Started example VM.");
 
     // Wait for VM to finish
diff --git a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
index 109c5e0..1d827b9 100644
--- a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
+++ b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
@@ -27,6 +27,8 @@
 import static com.google.common.truth.Truth.assertWithMessage;
 import static com.google.common.truth.TruthJUnit.assume;
 
+import static org.junit.Assume.assumeTrue;
+
 import android.app.Application;
 import android.app.Instrumentation;
 import android.content.ComponentCallbacks2;
@@ -69,13 +71,16 @@
 import java.io.Writer;
 import java.nio.file.Files;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.OptionalLong;
+import java.util.Random;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 
@@ -927,4 +932,169 @@
         }
         reportMetrics(requestAttestationTime, "request_attestation_time", "microsecond");
     }
+
+    List<Double> rpDataAccessWithExistingSession(boolean measureWrite) throws Exception {
+        assumeTrue(
+                "Rollback protected secrets are only available in Updatable VMs",
+                isUpdatableVmSupported());
+        final int NUM_WARMUPS = 10;
+        final int NUM_REQUESTS = 10_000;
+
+        VirtualMachineConfig config =
+                newVmConfigBuilderWithPayloadBinary("MicrodroidTestNativeLib.so")
+                        .setDebugLevel(DEBUG_LEVEL_NONE)
+                        .build();
+
+        byte[] data = new byte[32];
+        Arrays.fill(data, (byte) 0xcc);
+
+        List<Double> requestLatencies = new ArrayList<>(NUM_REQUESTS);
+        VirtualMachine vm = forceCreateNewVirtualMachine("rp_data_access", config);
+        TestResults testResult =
+                runVmTestService(
+                        TAG,
+                        vm,
+                        (ts, tr) -> {
+                            tr.mTimings = new long[NUM_REQUESTS];
+                            for (int i = 0; i < NUM_WARMUPS; i++) {
+                                ts.insecurelyWritePayloadRpData(data);
+                                ts.insecurelyReadPayloadRpData();
+                            }
+                            for (int i = 0; i < NUM_REQUESTS; i++) {
+                                long start = System.nanoTime();
+                                if (measureWrite) {
+                                    ts.insecurelyWritePayloadRpData(data);
+                                    tr.mTimings[i] = System.nanoTime() - start;
+                                } else {
+                                    tr.mPayloadRpData = ts.insecurelyReadPayloadRpData();
+                                    tr.mTimings[i] = System.nanoTime() - start;
+                                    assertThat(tr.mPayloadRpData).isEqualTo(data);
+                                }
+                            }
+                        });
+        // Correctness check.
+        testResult.assertNoException();
+        for (long timings : testResult.mTimings) {
+            requestLatencies.add((double) timings / NANO_TO_MICRO);
+        }
+        return requestLatencies;
+    }
+
+    @Test
+    public void rpDataReadWithExistingSession() throws Exception {
+        reportMetrics(
+                rpDataAccessWithExistingSession(false),
+                "latency/readRollbackProtectedSecretWithExistingSession",
+                "us");
+    }
+
+    @Test
+    public void rpDataWriteWithExistingSession() throws Exception {
+        reportMetrics(
+                rpDataAccessWithExistingSession(true),
+                "latency/writeRollbackProtectedSecretWithExistingSession",
+                "us");
+    }
+
+    List<Double> rpDataAccessWithRefreshingSession(boolean measureWrite) throws Exception {
+        assumeTrue(
+                "Rollback protected secrets are only available in Updatable VMs",
+                isUpdatableVmSupported());
+        final long vmSize = minMemoryRequired();
+        final int numVMs = 8;
+        final int NUM_REQUESTS = 10;
+        final long availableMem = getAvailableMemory();
+
+        // Let's not use more than half of the available memory
+        assume().withMessage("Available memory (" + availableMem + " bytes) too small")
+                .that((numVMs * vmSize) <= (availableMem / 2))
+                .isTrue();
+
+        VirtualMachineConfig config =
+                newVmConfigBuilderWithPayloadBinary("MicrodroidTestNativeLib.so")
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
+                        .setMemoryBytes(vmSize)
+                        .build();
+
+        byte[] data = new byte[32];
+        Arrays.fill(data, (byte) 0xcc);
+
+        List<Double> requestLatencies = new ArrayList<>(numVMs * NUM_REQUESTS);
+        CompletableFuture<TestResults>[] resultFutureList = new CompletableFuture[numVMs];
+        RunTestsAgainstTestService testToRun =
+                (ts, tr) -> {
+                    tr.mTimings = new long[NUM_REQUESTS];
+                    // Warm up request!
+                    ts.insecurelyWritePayloadRpData(data);
+                    for (int j = 0; j < NUM_REQUESTS; j++) {
+                        // Sleep time between 2 requests.
+                        // Randomized
+                        // between 200ms-300ms.
+                        long rnd_sleep_time = (long) (200.0 + new Random().nextDouble() * 100);
+                        Thread.sleep(rnd_sleep_time); // Sleep
+                        long start = System.nanoTime();
+                        if (measureWrite) {
+                            // Write
+                            ts.insecurelyWritePayloadRpData(data);
+                            tr.mTimings[j] = System.nanoTime() - start;
+
+                        } else {
+                            tr.mPayloadRpData = ts.insecurelyReadPayloadRpData();
+                            tr.mTimings[j] = System.nanoTime() - start;
+                            assertThat(tr.mPayloadRpData).isEqualTo(data);
+                        }
+                    }
+                };
+        for (int i = 0; i < numVMs; i++) {
+            final VirtualMachine vm =
+                    forceCreateNewVirtualMachine("rp_data_access_refresh" + i, config);
+            resultFutureList[i] =
+                    CompletableFuture.supplyAsync(
+                            () -> {
+                                try {
+                                    TestResults testResult = runVmTestService(TAG, vm, testToRun);
+                                    // Correctness check.
+                                    testResult.assertNoException();
+                                    return testResult;
+                                } catch (Exception e) {
+                                    throw new CompletionException(e);
+                                }
+                            });
+        }
+
+        for (int i = 0; i < numVMs; i++) {
+            TestResults tr = resultFutureList[i].get();
+            tr.assertNoException();
+            for (long timings : tr.mTimings) {
+                requestLatencies.add((double) timings / NANO_TO_MICRO);
+            }
+        }
+        return requestLatencies;
+    }
+
+    // The following benchmark corresponds to cases when payload access rollback protected secret,
+    // but there is no existing session with Secretkeeper - which could be the case when several VMs
+    // are attempting to establish a connection.
+    //
+    // Implementation detail of the API in such scenario: Microdroid attempts to access the secret
+    // from Secretkeeper -> gets an error ("UnknownKeyId") -> Refreshes the session (this includes
+    // several call to AuthGraphKey Exchange HAL) -> retries access.
+    //
+    // Essentially this latency is (Failed Secretkeeper access from pVM + AuthGraphKeyExchange
+    // protocol between pVM & Secretkeeper + Successful Secretkeeper access from pVM)
+    @Test
+    public void rpDataReadWithRefreshingSession() throws Exception {
+        reportMetrics(
+                rpDataAccessWithRefreshingSession(false),
+                "latency/readRollbackProtectedSecretWithRefreshSession",
+                "us");
+    }
+
+    @Test
+    public void rpDataWriteWithRefreshingSession() throws Exception {
+        reportMetrics(
+                rpDataAccessWithRefreshingSession(true),
+                "latency/writeRollbackProtectedSecretWithRefreshSession",
+                "us");
+    }
 }
diff --git a/tests/early_vm_test/src/main.rs b/tests/early_vm_test/src/main.rs
index a3c80ca..7d630f8 100644
--- a/tests/early_vm_test/src/main.rs
+++ b/tests/early_vm_test/src/main.rs
@@ -96,7 +96,6 @@
         None, // console_out
         None, // log
         None, // dump_dt
-        None, // callback
     )
     .context("Failed to create VM")?;
 
diff --git a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
index c05fb0b..94f7ced 100644
--- a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
+++ b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
@@ -27,9 +27,11 @@
 import static org.junit.Assume.assumeFalse;
 import static org.junit.Assume.assumeTrue;
 
+import android.app.ActivityManager;
 import android.app.Instrumentation;
 import android.app.UiAutomation;
 import android.content.Context;
+import android.os.Build;
 import android.os.ParcelFileDescriptor;
 import android.os.SystemProperties;
 import android.system.Os;
@@ -79,6 +81,10 @@
                                     "microdroid_16k",
                                     "microdroid_gki-android15-6.6")));
 
+    private static final long ONE_MEBI = 1024 * 1024;
+    private static final long MIN_MEM_ARM64 = 170 * ONE_MEBI;
+    private static final long MIN_MEM_X86_64 = 196 * ONE_MEBI;
+
     public static boolean isCuttlefish() {
         return getDeviceProperties().isCuttlefish();
     }
@@ -393,6 +399,10 @@
             return mProcessedBootTimeMetrics;
         }
 
+        // Stopping a virtual machine is like pulling the plug on a real computer. VM may be left in
+        // an inconsistent state.
+        // For a graceful shutdown, request the payload to call {@code exit()} and wait for
+        // VirtualMachineCallback#onPayloadFinished} to be called.
         protected void forceStop(VirtualMachine vm) {
             try {
                 vm.stop();
@@ -722,7 +732,6 @@
                     public void onPayloadFinished(VirtualMachine vm, int exitCode) {
                         Log.i(logTag, "onPayloadFinished: " + exitCode);
                         payloadFinished.complete(true);
-                        forceStop(vm);
                     }
                 };
 
@@ -733,6 +742,26 @@
         return testResults;
     }
 
+    protected long getAvailableMemory() {
+        ActivityManager am = getContext().getSystemService(ActivityManager.class);
+        ActivityManager.MemoryInfo memoryInfo = new ActivityManager.MemoryInfo();
+        am.getMemoryInfo(memoryInfo);
+        return memoryInfo.availMem;
+    }
+
+    protected long minMemoryRequired() {
+        assertThat(Build.SUPPORTED_ABIS).isNotEmpty();
+        String primaryAbi = Build.SUPPORTED_ABIS[0];
+        switch (primaryAbi) {
+            case "x86_64":
+                return MIN_MEM_X86_64;
+            case "arm64-v8a":
+            case "arm64-v8a-hwasan":
+                return MIN_MEM_ARM64;
+        }
+        throw new AssertionError("Unsupported ABI: " + primaryAbi);
+    }
+
     @FunctionalInterface
     protected interface RunTestsAgainstTestService {
         void runTests(ITestService testService, TestResults testResults) throws Exception;
diff --git a/tests/old_images_avf_test/src/main.rs b/tests/old_images_avf_test/src/main.rs
index 018a80e..d3ab5eb 100644
--- a/tests/old_images_avf_test/src/main.rs
+++ b/tests/old_images_avf_test/src/main.rs
@@ -92,7 +92,7 @@
         AVirtualMachineRawConfig_setMemoryMiB(config, VM_MEMORY_MB);
     }
 
-    let mut vm = std::ptr::null_mut();
+    let mut vm = std::ptr::null();
     let mut service = std::ptr::null_mut();
 
     ensure!(
@@ -132,7 +132,7 @@
 
     // SAFETY: vm is the only reference to a valid object
     unsafe {
-        AVirtualMachine_start(vm);
+        AVirtualMachine_start(vm, None /* callback */);
     }
 
     info!("VM started");
diff --git a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
index a2b4747..7431a72 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -32,14 +32,12 @@
 import static com.google.common.truth.TruthJUnit.assume;
 
 import static org.junit.Assert.assertThrows;
-import static org.junit.Assert.assertTrue;
 import static org.junit.Assume.assumeFalse;
 import static org.junit.Assume.assumeTrue;
 
 import static java.nio.file.StandardCopyOption.REPLACE_EXISTING;
 import static java.util.stream.Collectors.toList;
 
-import android.app.ActivityManager;
 import android.app.Instrumentation;
 import android.app.UiAutomation;
 import android.content.ComponentName;
@@ -47,7 +45,6 @@
 import android.content.ContextWrapper;
 import android.content.Intent;
 import android.content.ServiceConnection;
-import android.os.Build;
 import android.os.IBinder;
 import android.os.Parcel;
 import android.os.ParcelFileDescriptor;
@@ -119,6 +116,7 @@
 import java.util.OptionalLong;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
@@ -171,11 +169,6 @@
     public void tearDown() {
         revokePermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION);
     }
-
-    private static final long ONE_MEBI = 1024 * 1024;
-
-    private static final long MIN_MEM_ARM64 = 170 * ONE_MEBI;
-    private static final long MIN_MEM_X86_64 = 196 * ONE_MEBI;
     private static final String EXAMPLE_STRING = "Literally any string!! :)";
 
     private static final String VM_SHARE_APP_PACKAGE_NAME = "com.android.microdroid.vmshare_app";
@@ -1875,6 +1868,7 @@
     }
 
     @Test
+    @CddTest
     public void rollbackProtectedDataOfPayload() throws Exception {
         assumeSupportedDevice();
         // Rollback protected data is only possible if Updatable VMs is supported -
@@ -1900,7 +1894,7 @@
                         (ts, tr) -> {
                             tr.mPayloadRpData = ts.insecurelyReadPayloadRpData();
                         });
-        // ainsecurelyReadPayloadRpData()` must've failed since no data was ever written!
+        // `insecurelyReadPayloadRpData()` must've failed since no data was ever written!
         assertWithMessage("The read (unexpectedly) succeeded!")
                 .that(testResults.mException)
                 .isNotNull();
@@ -1931,6 +1925,62 @@
     }
 
     @Test
+    public void rollbackProtectedDataCanBeAccessedPostConnectionExpiration() throws Exception {
+        final long vmSize = minMemoryRequired();
+        // The reference implementation of Secretkeeper maintains 4 live session keys,
+        // dropping the oldest one when new connections are requested. Therefore we spin 8 VMs
+        // asynchronously.
+        // Within a VM, wait for 5 sec (> Microdroid boot time) and trigger rp data access
+        // hoping at least some of the connection between VM <-> Secretkeeper are expired.
+        final int numVMs = 8;
+        final long availableMem = getAvailableMemory();
+
+        // Let's not use more than half of the available memory
+        assume().withMessage("Available memory (" + availableMem + " bytes) too small")
+                .that((numVMs * vmSize) <= (availableMem / 2))
+                .isTrue();
+
+        VirtualMachineConfig config =
+                newVmConfigBuilderWithPayloadBinary("MicrodroidTestNativeLib.so")
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
+                        .setMemoryBytes(vmSize)
+                        .build();
+        byte[] data = new byte[32];
+        Arrays.fill(data, (byte) 0xcc);
+
+        CompletableFuture<TestResults>[] resultFutureList = new CompletableFuture[numVMs];
+        for (int i = 0; i < numVMs; i++) {
+            final VirtualMachine vm =
+                    forceCreateNewVirtualMachine("test_sk_session_expiration_vm_" + i, config);
+            resultFutureList[i] =
+                    CompletableFuture.supplyAsync(
+                            () -> {
+                                try {
+                                    TestResults testResults =
+                                            runVmTestService(
+                                                    TAG,
+                                                    vm,
+                                                    (ts, tr) -> {
+                                                        ts.insecurelyWritePayloadRpData(data);
+                                                        Thread.sleep(5 * 1000); // 5 seconds of wait
+                                                        tr.mPayloadRpData =
+                                                                ts.insecurelyReadPayloadRpData();
+                                                    });
+                                    return testResults;
+                                } catch (Exception e) {
+                                    throw new CompletionException(e);
+                                }
+                            });
+        }
+
+        for (int i = 0; i < numVMs; i++) {
+            TestResults testResult = resultFutureList[i].get();
+            testResult.assertNoException();
+            assertThat(testResult.mPayloadRpData).isEqualTo(data);
+        }
+    }
+
+    @Test
     @CddTest
     public void isNewInstanceTest() throws Exception {
         assumeSupportedDevice();
@@ -2464,25 +2514,7 @@
     }
 
     @Test
-    public void kernelVersionRequirement() throws Exception {
-        assumeVsrCompliant();
-        int firstApiLevel = SystemProperties.getInt("ro.product.first_api_level", 0);
-        assume().withMessage("Skip on devices launched before Android 14 (API level 34)")
-                .that(firstApiLevel)
-                .isAtLeast(34);
-
-        String[] tokens = KERNEL_VERSION.split("\\.");
-        int major = Integer.parseInt(tokens[0]);
-        int minor = Integer.parseInt(tokens[1]);
-
-        // Check kernel version >= 5.15
-        assertTrue(major >= 5);
-        if (major == 5) {
-            assertTrue(minor >= 15);
-        }
-    }
-
-    @Test
+    @CddTest
     public void createAndRunRustVm() throws Exception {
         // This test is here mostly to exercise the Rust wrapper around the VM Payload API.
         // We're testing the same functionality as in other tests, the only difference is
@@ -2686,6 +2718,7 @@
     }
 
     @Test
+    @GmsTest(requirements = {"GMS-3-7.1-001.002"})
     public void pageSize() throws Exception {
         assumeSupportedDevice();
 
@@ -2771,13 +2804,6 @@
         }
     }
 
-    private long getAvailableMemory() {
-        ActivityManager am = getContext().getSystemService(ActivityManager.class);
-        ActivityManager.MemoryInfo memoryInfo = new ActivityManager.MemoryInfo();
-        am.getMemoryInfo(memoryInfo);
-        return memoryInfo.availMem;
-    }
-
     private VirtualMachineDescriptor toParcelFromParcel(VirtualMachineDescriptor descriptor) {
         Parcel parcel = Parcel.obtain();
         descriptor.writeToParcel(parcel, 0);
@@ -2810,17 +2836,4 @@
         Exception e = assertThrows(VirtualMachineException.class, runnable);
         assertThat(e).hasMessageThat().contains(expectedContents);
     }
-
-    private long minMemoryRequired() {
-        assertThat(Build.SUPPORTED_ABIS).isNotEmpty();
-        String primaryAbi = Build.SUPPORTED_ABIS[0];
-        switch (primaryAbi) {
-            case "x86_64":
-                return MIN_MEM_X86_64;
-            case "arm64-v8a":
-            case "arm64-v8a-hwasan":
-                return MIN_MEM_ARM64;
-        }
-        throw new AssertionError("Unsupported ABI: " + primaryAbi);
-    }
 }
diff --git a/tests/vm_accessor/accessor/src/run.rs b/tests/vm_accessor/accessor/src/run.rs
index 6dcc507..5bdb8f1 100644
--- a/tests/vm_accessor/accessor/src/run.rs
+++ b/tests/vm_accessor/accessor/src/run.rs
@@ -129,10 +129,9 @@
         None,                    /* console_in */
         Some(android_log_fd()?), /* log */
         None,                    /* dump_dt */
-        Some(Box::new(Callback {})),
     )
     .context("Failed to create VM")?;
-    vm.start().context("Failed to start VM")?;
+    vm.start(Some(Box::new(Callback {}))).context("Failed to start VM")?;
 
     info!("started IAccessor VM with CID {}", vm.cid());
 
diff --git a/tests/vmbase_example/src/main.rs b/tests/vmbase_example/src/main.rs
index cbe90d8..d427164 100644
--- a/tests/vmbase_example/src/main.rs
+++ b/tests/vmbase_example/src/main.rs
@@ -115,10 +115,9 @@
         /* consoleIn */ None,
         Some(log_writer),
         /* dump_dt */ None,
-        None,
     )
     .context("Failed to create VM")?;
-    vm.start().context("Failed to start VM")?;
+    vm.start(None).context("Failed to start VM")?;
     info!("Started example VM.");
 
     // Wait for VM to finish, and check that it shut down cleanly.
diff --git a/tests/vts/src/vts_libavf_test.rs b/tests/vts/src/vts_libavf_test.rs
index dc37aad..c59271a 100644
--- a/tests/vts/src/vts_libavf_test.rs
+++ b/tests/vts/src/vts_libavf_test.rs
@@ -19,6 +19,8 @@
 use std::fs::File;
 use std::io::{self, BufWriter, Write};
 use std::os::fd::IntoRawFd;
+use std::sync::mpsc::{self, Sender};
+use std::sync::{LazyLock, Mutex};
 use std::time::{Duration, Instant};
 use vsock::{VsockListener, VsockStream, VMADDR_CID_HOST};
 
@@ -33,6 +35,13 @@
 const WRITE_TIMEOUT: Duration = Duration::from_secs(10);
 const STOP_TIMEOUT: timespec = timespec { tv_sec: 10, tv_nsec: 0 };
 
+static ON_STOPPED_EVENT: LazyLock<Mutex<Sender<(usize, AVirtualMachineStopReason)>>> =
+    LazyLock::new(|| {
+        // Returning stub here because `Receiver` isn't `Send`.
+        let (tx, _) = mpsc::channel();
+        Mutex::new(tx)
+    });
+
 /// Processes the request in the service VM.
 fn process_request(vsock_stream: &mut VsockStream, request: Request) -> Result<Response> {
     write_request(vsock_stream, &ServiceVmRequest::Process(request))?;
@@ -73,11 +82,23 @@
     }
 }
 
+unsafe extern "C" fn on_stopped(vm: *const AVirtualMachine, reason: AVirtualMachineStopReason) {
+    ON_STOPPED_EVENT.lock().unwrap().send((vm as usize, reason)).unwrap();
+
+    // SAFETY: `vm` is a valid pointer created by AVirtualMachine_create().
+    unsafe {
+        AVirtualMachine_destroy(vm);
+    }
+}
+
 fn run_rialto(protected_vm: bool) -> Result<()> {
     let kernel_file =
         File::open("/data/local/tmp/rialto.bin").context("Failed to open kernel file")?;
     let kernel_fd = kernel_file.into_raw_fd();
 
+    let (tx, rx) = mpsc::channel();
+    (*ON_STOPPED_EVENT.lock().unwrap()) = tx;
+
     // SAFETY: AVirtualMachineRawConfig_create() isn't unsafe but rust_bindgen forces it to be seen
     // as unsafe
     let config = unsafe { AVirtualMachineRawConfig_create() };
@@ -92,7 +113,7 @@
         AVirtualMachineRawConfig_setMemoryMiB(config, VM_MEMORY_MB);
     }
 
-    let mut vm = std::ptr::null_mut();
+    let mut vm = std::ptr::null();
     let mut service = std::ptr::null_mut();
 
     ensure!(
@@ -101,11 +122,6 @@
         "AVirtualizationService_create failed"
     );
 
-    scopeguard::defer! {
-        // SAFETY: service is a valid pointer to AVirtualizationService
-        unsafe { AVirtualizationService_destroy(service); }
-    }
-
     ensure!(
         // SAFETY: &mut vm is a valid pointer to *AVirtualMachine
         unsafe {
@@ -132,9 +148,11 @@
 
     // SAFETY: vm is the only reference to a valid object
     unsafe {
-        AVirtualMachine_start(vm);
+        AVirtualMachine_start(vm, Some(on_stopped));
     }
 
+    let vm_ptr = vm as usize;
+
     info!("VM started");
 
     let mut vsock_stream = listener_thread.join().unwrap()?;
@@ -166,6 +184,14 @@
         "AVirtualMachine_waitForStop failed"
     );
 
+    assert_eq!(AVirtualMachineStopReason::AVIRTUAL_MACHINE_SHUTDOWN, stop_reason);
+
+    let timeout = Duration::from_secs(STOP_TIMEOUT.tv_sec.try_into().unwrap());
+    let (stopped_vm_ptr, stopped_reason) =
+        rx.recv_timeout(timeout).expect("Callback should have been called");
+    assert_eq!(stopped_vm_ptr, vm_ptr);
+    assert_eq!(stopped_reason, stop_reason);
+
     info!("stopped");
 
     Ok(())