diff --git a/virtualizationservice/Android.bp b/virtualizationservice/Android.bp
index 26d41c9..d6f4607 100644
--- a/virtualizationservice/Android.bp
+++ b/virtualizationservice/Android.bp
@@ -71,5 +71,8 @@
 rust_test {
     name: "virtualizationservice_device_test",
     defaults: ["virtualizationservice_defaults"],
+    rustlibs: [
+        "libtempfile",
+    ],
     test_suites: ["general-tests"],
 }
diff --git a/virtualizationservice/src/payload.rs b/virtualizationservice/src/payload.rs
index 4190fbb..233e74b 100644
--- a/virtualizationservice/src/payload.rs
+++ b/virtualizationservice/src/payload.rs
@@ -27,7 +27,9 @@
 use microdroid_metadata::{ApexPayload, ApkPayload, Metadata, PayloadConfig, PayloadMetadata};
 use microdroid_payload_config::{ApexConfig, VmPayloadConfig};
 use once_cell::sync::OnceCell;
-use packagemanager_aidl::aidl::android::content::pm::IPackageManagerNative::IPackageManagerNative;
+use packagemanager_aidl::aidl::android::content::pm::{
+    IPackageManagerNative::IPackageManagerNative, StagedApexInfo::StagedApexInfo,
+};
 use regex::Regex;
 use serde::Deserialize;
 use serde_xml_rs::from_reader;
@@ -48,7 +50,7 @@
 const PACKAGE_MANAGER_NATIVE_SERVICE: &str = "package_native";
 
 /// Represents the list of APEXes
-#[derive(Clone, Debug, Deserialize)]
+#[derive(Clone, Debug, Deserialize, Eq, PartialEq)]
 struct ApexInfoList {
     #[serde(rename = "apex-info")]
     list: Vec<ApexInfo>,
@@ -58,6 +60,8 @@
 struct ApexInfo {
     #[serde(rename = "moduleName")]
     name: String,
+    #[serde(rename = "versionCode")]
+    version: u64,
     #[serde(rename = "modulePath")]
     path: PathBuf,
 
@@ -101,6 +105,40 @@
             Ok(apex_info_list)
         })
     }
+
+    // Override apex info with the staged one
+    fn override_staged_apex(&mut self, staged_apex_info: &StagedApexInfo) -> Result<()> {
+        let mut need_to_add: Option<ApexInfo> = None;
+        for apex_info in self.list.iter_mut() {
+            if staged_apex_info.moduleName == apex_info.name {
+                if apex_info.is_active && apex_info.is_factory {
+                    // Copy the entry to the end as factory/non-active after the loop
+                    // to keep the factory version. Typically this step is unncessary,
+                    // but some apexes (like sharedlibs) need to be kept even if it's inactive.
+                    need_to_add.replace(ApexInfo { is_active: false, ..apex_info.clone() });
+                    // And make this one as non-factory. Note that this one is still active
+                    // and overridden right below.
+                    apex_info.is_factory = false;
+                }
+                // Active one is overridden with the staged one.
+                if apex_info.is_active {
+                    apex_info.version = staged_apex_info.versionCode as u64;
+                    apex_info.path = PathBuf::from(&staged_apex_info.diskImagePath);
+                    apex_info.has_classpath_jar = staged_apex_info.hasClassPathJars;
+                    apex_info.last_update_seconds = last_updated(&apex_info.path)?;
+                }
+            }
+        }
+        if let Some(info) = need_to_add {
+            self.list.push(info);
+        }
+        Ok(())
+    }
+}
+
+fn last_updated<P: AsRef<Path>>(path: P) -> Result<u64> {
+    let metadata = metadata(path)?;
+    Ok(metadata.modified()?.duration_since(SystemTime::UNIX_EPOCH)?.as_secs())
 }
 
 impl ApexInfo {
@@ -137,19 +175,11 @@
                     .context("Failed to get service when prefer_staged is set.")?;
             let staged =
                 pm.getStagedApexModuleNames().context("getStagedApexModuleNames failed")?;
-            for apex_info in list.list.iter_mut() {
-                if staged.contains(&apex_info.name) {
-                    if let Some(staged_apex_info) =
-                        pm.getStagedApexInfo(&apex_info.name).context("getStagedApexInfo failed")?
-                    {
-                        apex_info.path = PathBuf::from(staged_apex_info.diskImagePath);
-                        apex_info.has_classpath_jar = staged_apex_info.hasClassPathJars;
-                        let metadata = metadata(&apex_info.path)?;
-                        apex_info.last_update_seconds =
-                            metadata.modified()?.duration_since(SystemTime::UNIX_EPOCH)?.as_secs();
-                        // by definition, staged apex can't be a factory apex.
-                        apex_info.is_factory = false;
-                    }
+            for name in staged {
+                if let Some(staged_apex_info) =
+                    pm.getStagedApexInfo(&name).context("getStagedApexInfo failed")?
+                {
+                    list.override_staged_apex(&staged_apex_info)?;
                 }
             }
         }
@@ -243,8 +273,13 @@
     let apex_list = pm.get_apex_list(vm_payload_config.prefer_staged)?;
 
     // collect APEXes from config
-    let apex_infos =
+    let mut apex_infos =
         collect_apex_infos(&apex_list, &vm_payload_config.apexes, app_config.debugLevel);
+
+    // Pass sorted list of apexes. Sorting key shouldn't use `path` because it will change after
+    // reboot with prefer_staged. `last_update_seconds` is added to distinguish "samegrade"
+    // update.
+    apex_infos.sort_by_key(|info| (&info.name, &info.version, &info.last_update_seconds));
     info!("Microdroid payload APEXes: {:?}", apex_infos.iter().map(|ai| &ai.name));
 
     let metadata_file = make_metadata_file(app_config, &apex_infos, temporary_directory)?;
@@ -422,6 +457,7 @@
 #[cfg(test)]
 mod tests {
     use super::*;
+    use tempfile::NamedTempFile;
 
     #[test]
     fn test_find_apex_names_in_classpath() {
@@ -560,4 +596,90 @@
             ]
         );
     }
+
+    #[test]
+    fn test_prefer_staged_apex_with_factory_active_apex() {
+        let single_apex = ApexInfo {
+            name: "foo".to_string(),
+            version: 1,
+            path: PathBuf::from("foo.apex"),
+            is_factory: true,
+            is_active: true,
+            ..Default::default()
+        };
+        let mut apex_info_list = ApexInfoList { list: vec![single_apex.clone()] };
+
+        let staged = NamedTempFile::new().unwrap();
+        apex_info_list
+            .override_staged_apex(&StagedApexInfo {
+                moduleName: "foo".to_string(),
+                versionCode: 2,
+                diskImagePath: staged.path().to_string_lossy().to_string(),
+                ..Default::default()
+            })
+            .expect("should be ok");
+
+        assert_eq!(
+            apex_info_list,
+            ApexInfoList {
+                list: vec![
+                    ApexInfo {
+                        version: 2,
+                        is_factory: false,
+                        path: staged.path().to_owned(),
+                        last_update_seconds: last_updated(staged.path()).unwrap(),
+                        ..single_apex.clone()
+                    },
+                    ApexInfo { is_active: false, ..single_apex },
+                ],
+            }
+        );
+    }
+
+    #[test]
+    fn test_prefer_staged_apex_with_factory_and_inactive_apex() {
+        let factory_apex = ApexInfo {
+            name: "foo".to_string(),
+            version: 1,
+            path: PathBuf::from("foo.apex"),
+            is_factory: true,
+            ..Default::default()
+        };
+        let active_apex = ApexInfo {
+            name: "foo".to_string(),
+            version: 2,
+            path: PathBuf::from("foo.downloaded.apex"),
+            is_active: true,
+            ..Default::default()
+        };
+        let mut apex_info_list =
+            ApexInfoList { list: vec![factory_apex.clone(), active_apex.clone()] };
+
+        let staged = NamedTempFile::new().unwrap();
+        apex_info_list
+            .override_staged_apex(&StagedApexInfo {
+                moduleName: "foo".to_string(),
+                versionCode: 3,
+                diskImagePath: staged.path().to_string_lossy().to_string(),
+                ..Default::default()
+            })
+            .expect("should be ok");
+
+        assert_eq!(
+            apex_info_list,
+            ApexInfoList {
+                list: vec![
+                    // factory apex isn't touched
+                    factory_apex,
+                    // update active one
+                    ApexInfo {
+                        version: 3,
+                        path: staged.path().to_owned(),
+                        last_update_seconds: last_updated(staged.path()).unwrap(),
+                        ..active_apex
+                    },
+                ],
+            }
+        );
+    }
 }