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..0c79f5f 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>,
@@ -101,6 +103,39 @@
             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.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 +172,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)?;
                 }
             }
         }
@@ -422,6 +449,7 @@
 #[cfg(test)]
 mod tests {
     use super::*;
+    use tempfile::NamedTempFile;
 
     #[test]
     fn test_find_apex_names_in_classpath() {
@@ -560,4 +588,83 @@
             ]
         );
     }
+
+    #[test]
+    fn test_prefer_staged_apex_with_factory_active_apex() {
+        let single_apex = ApexInfo {
+            name: "foo".to_string(),
+            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(),
+                diskImagePath: staged.path().to_string_lossy().to_string(),
+                ..Default::default()
+            })
+            .expect("should be ok");
+
+        assert_eq!(
+            apex_info_list,
+            ApexInfoList {
+                list: vec![
+                    ApexInfo {
+                        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(),
+            path: PathBuf::from("foo.apex"),
+            is_factory: true,
+            ..Default::default()
+        };
+        let active_apex = ApexInfo {
+            name: "foo".to_string(),
+            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(),
+                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 {
+                        path: staged.path().to_owned(),
+                        last_update_seconds: last_updated(staged.path()).unwrap(),
+                        ..active_apex
+                    },
+                ],
+            }
+        );
+    }
 }