From c148b7df6c53f04e86b1daefd6598e5203b98611 Mon Sep 17 00:00:00 2001
From: Elie Kheirallah <khei@google.com>
Date: Fri, 4 Sep 2026 00:38:31 +0000
Subject: [PATCH] assemble_cvd: Skip rebuilding OS composite disk on snapshot
 restore

When restoring from a snapshot, the base images and instance disks
have already been restored from the snapshot directory by
RestoreHostFiles(). Checking filesystem timestamps to determine whether
to rebuild composite disks during snapshot restore can lead to
creating_os_disk becoming true, which causes assemble_cvd to abort
with 'Restoring from snapshot requires not creating OS disks'.

Only evaluate WillRebuildCompositeDisk() and set creating_os_disk when
snapshot_path is empty.

Bug: 554545294
---
 .../commands/assemble_cvd/assemble_cvd.cc     | 72 ++++++++++---------
 1 file changed, 38 insertions(+), 34 deletions(-)

diff --git a/base/cvd/cuttlefish/host/commands/assemble_cvd/assemble_cvd.cc b/base/cvd/cuttlefish/host/commands/assemble_cvd/assemble_cvd.cc
index 6404b1cc2c4..e51a852c253 100644
--- a/base/cvd/cuttlefish/host/commands/assemble_cvd/assemble_cvd.cc
+++ b/base/cvd/cuttlefish/host/commands/assemble_cvd/assemble_cvd.cc
@@ -376,45 +376,49 @@ Result<const CuttlefishConfig*> InitFilesystemAndCreateConfig(
     // take the max value of modem_simulator_instance_number in each instance
     // which is used for preserving/deleting iccprofile_for_simX.xml files
     int modem_simulator_count = 0;
-
-    bool creating_os_disk = false;
-    // if any device needs to rebuild its composite disk,
-    // then don't preserve any files and delete everything.
-
-    std::vector<std::vector<std::unique_ptr<ImageFile>>> image_files =
-        InstanceImageFiles(config, boot_image);
-
-    size_t index = 0;
     for (const auto& instance : config.Instances()) {
-      CF_EXPECT_LE(index, image_files.size());
-      const std::vector<std::unique_ptr<ImageFile>>& instance_image_files =
-          image_files[index];
-
-      std::optional<ChromeOsStateImage> chrome_os_state =
-          CF_EXPECT(ChromeOsStateImage::Reuse(instance));
-      Result<DiskBuilder> os_builder = OsCompositeDiskBuilder(
-          config, instance, chrome_os_state, instance_image_files,
-          android_builds.ForIndex(index), system_image_dir);
-      if (!os_builder.has_value()) {
-        creating_os_disk = true;
-      } else {
-        creating_os_disk |= CF_EXPECT(os_builder->WillRebuildCompositeDisk());
-      }
-      if (instance.ap_boot_flow() != APBootFlow::None) {
-        auto ap_builder = ApCompositeDiskBuilder(config, instance);
-        creating_os_disk |= CF_EXPECT(ap_builder.WillRebuildCompositeDisk());
-      }
       if (instance.modem_simulator_instance_number() > modem_simulator_count) {
         modem_simulator_count = instance.modem_simulator_instance_number();
       }
-      index++;
     }
-    // TODO(schuffelen): Add smarter decision for when to delete runtime files.
-    // Files like NVChip are tightly bound to Android keymint and should be
-    // deleted when userdata is reset. However if the user has ever run without
-    // the overlay, then we want to keep this until userdata.img was externally
-    // replaced.
-    creating_os_disk &= FLAGS_use_overlay;
+
+    bool creating_os_disk = false;
+    if (snapshot_path.empty()) {
+      // if any device needs to rebuild its composite disk,
+      // then don't preserve any files and delete everything.
+
+      std::vector<std::vector<std::unique_ptr<ImageFile>>> image_files =
+          InstanceImageFiles(config, boot_image);
+
+      size_t index = 0;
+      for (const auto& instance : config.Instances()) {
+        CF_EXPECT_LE(index, image_files.size());
+        const std::vector<std::unique_ptr<ImageFile>>& instance_image_files =
+            image_files[index];
+
+        std::optional<ChromeOsStateImage> chrome_os_state =
+            CF_EXPECT(ChromeOsStateImage::Reuse(instance));
+        Result<DiskBuilder> os_builder = OsCompositeDiskBuilder(
+            config, instance, chrome_os_state, instance_image_files,
+            android_builds.ForIndex(index), system_image_dir);
+        if (!os_builder.has_value()) {
+          creating_os_disk = true;
+        } else {
+          creating_os_disk |= CF_EXPECT(os_builder->WillRebuildCompositeDisk());
+        }
+        if (instance.ap_boot_flow() != APBootFlow::None) {
+          auto ap_builder = ApCompositeDiskBuilder(config, instance);
+          creating_os_disk |= CF_EXPECT(ap_builder.WillRebuildCompositeDisk());
+        }
+        index++;
+      }
+      // TODO(schuffelen): Add smarter decision for when to delete runtime files.
+      // Files like NVChip are tightly bound to Android keymint and should be
+      // deleted when userdata is reset. However if the user has ever run without
+      // the overlay, then we want to keep this until userdata.img was externally
+      // replaced.
+      creating_os_disk &= FLAGS_use_overlay;
+    }
 
     std::set<std::string> preserving =
         CF_EXPECT(PreservingOnResume(creating_os_disk, modem_simulator_count),