Eliminate rootfs copy: populate directly into target filesystem

Restructure the QCOW2 build pipeline so the target filesystem is created
and mounted before Phase 1, allowing rootfs population directly into it.
This removes the wasteful and fragile copy_rootfs step that was the source
of merged-/usr symlink breakage and fixes the ZFS shadow-mount bug where
the BE dataset was mounted over the staging root.

New flow: Phase 2 prepare → Phase 1 populate → Phase 2 finalize → cleanup.
OCI/Artifact targets are unchanged.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

crates/forge-engine/src/lib.rs

@@ -35,19 +35,24 @@ impl<'a> BuildContext<'a> {
         for target in targets {
             info!(target = %target.name, kind = %target.kind, "Building target");
 
-            // Phase 1: Assemble rootfs
+            match target.kind {
+                TargetKind::Qcow2 => {
+                    self.build_qcow2(target).await?;
+                }
+                _ => {
+                    // OCI/Artifact: Phase 1 creates its own staging dir, Phase 2 packs from it
                     let phase1_result =
                         phase1::execute(self.spec, self.files_dir, self.runner).await?;
 
-            // Phase 2: Produce target artifact
                     phase2::execute(
                         target,
                         &phase1_result.staging_root,
                         self.files_dir,
                         self.output_dir,
-                self.runner,
                     )
                     .await?;
+                }
+            }
 
             info!(target = %target.name, "Target built successfully");
         }
@@ -55,6 +60,40 @@ impl<'a> BuildContext<'a> {
         Ok(())
     }
 
+    /// QCOW2 build: prepare disk → populate rootfs directly → finalize → cleanup.
+    async fn build_qcow2(&self, target: &Target) -> Result<(), ForgeError> {
+        // Phase 2 prepare: create disk image, partition/format, mount target filesystem
+        let prepared =
+            phase2::qcow2::prepare_qcow2(target, self.output_dir, self.runner).await?;
+
+        // Phase 1: populate directly into the mounted target filesystem
+        let phase1_result =
+            phase1::execute_into(self.spec, self.files_dir, self.runner, prepared.root_mount())
+                .await;
+
+        // Phase 2 finalize: install bootloader, unmount (only if Phase 1 succeeded)
+        let finalize_result = if phase1_result.is_ok() {
+            phase2::qcow2::finalize_qcow2(&prepared, self.runner).await
+        } else {
+            Ok(())
+        };
+
+        // Cleanup always runs: detach loopback, convert if everything succeeded
+        let convert = phase1_result.is_ok() && finalize_result.is_ok();
+        let cleanup_result =
+            phase2::qcow2::cleanup_qcow2(prepared, convert, self.runner).await;
+
+        // Propagate errors in order of occurrence
+        phase1_result?;
+        finalize_result?;
+        cleanup_result?;
+
+        // Auto-push to OCI registry if configured
+        phase2::push_qcow2_if_configured(target, self.output_dir).await?;
+
+        Ok(())
+    }
+
     fn select_targets(&self, target_name: Option<&str>) -> Result<Vec<&Target>, ForgeError> {
         match target_name {
             Some(name) => {

crates/forge-engine/src/phase1/mod.rs

@@ -21,45 +21,57 @@ pub struct Phase1Result {
     pub _staging_dir: tempfile::TempDir,
 }
 
-/// Execute Phase 1: assemble a rootfs in a staging directory from the spec.
+/// Populate a rootfs into an existing directory (for QCOW2 targets that mount
+/// the target filesystem first).
+pub async fn execute_into(
+    spec: &ImageSpec,
+    files_dir: &Path,
+    runner: &dyn ToolRunner,
+    root: &Path,
+) -> Result<(), ForgeError> {
+    let distro = DistroFamily::from_distro_str(spec.distro.as_deref());
+    info!(name = %spec.metadata.name, ?distro, "Starting Phase 1: rootfs assembly");
+
+    let root_str = root.to_str().unwrap();
+    info!(root = root_str, "Populating rootfs");
+
+    // 1. Extract base tarball
+    if let Some(ref base) = spec.base {
+        staging::extract_base_tarball(base, &root.to_path_buf())?;
+    }
+
+    // 2. Distro-specific package management
+    match distro {
+        DistroFamily::OmniOS => execute_ips(spec, root_str, files_dir, runner).await?,
+        DistroFamily::Ubuntu => execute_apt(spec, root_str, runner).await?,
+    }
+
+    // 3. Apply customizations (common)
+    for customization in &spec.customizations {
+        customizations::apply(customization, root)?;
+    }
+
+    // 4. Apply overlays (common)
+    for overlay_block in &spec.overlays {
+        overlays::apply_overlays(&overlay_block.actions, root, files_dir, runner).await?;
+    }
+
+    info!("Phase 1 complete: rootfs assembled");
+    Ok(())
+}
+
+/// Execute Phase 1: create a staging directory and assemble a rootfs from the spec.
 ///
-/// Dispatches to the appropriate distro-specific path based on the `distro` field,
+/// For OCI/Artifact targets that manage their own staging directory.
-/// then applies common customizations and overlays.
+/// Delegates to `execute_into()` after creating the tempdir.
 pub async fn execute(
     spec: &ImageSpec,
     files_dir: &Path,
     runner: &dyn ToolRunner,
 ) -> Result<Phase1Result, ForgeError> {
-    let distro = DistroFamily::from_distro_str(spec.distro.as_deref());
-    info!(name = %spec.metadata.name, ?distro, "Starting Phase 1: rootfs assembly");
-
-    // 1. Create staging directory
     let (staging_dir, staging_root) = staging::create_staging()?;
-    let root = staging_root.to_str().unwrap();
-    info!(root, "Staging directory created");
 
-    // 2. Extract base tarball
+    execute_into(spec, files_dir, runner, &staging_root).await?;
-    if let Some(ref base) = spec.base {
-        staging::extract_base_tarball(base, &staging_root)?;
-    }
-
-    // 3. Distro-specific package management
-    match distro {
-        DistroFamily::OmniOS => execute_ips(spec, root, files_dir, runner).await?,
-        DistroFamily::Ubuntu => execute_apt(spec, root, runner).await?,
-    }
-
-    // 4. Apply customizations (common)
-    for customization in &spec.customizations {
-        customizations::apply(customization, &staging_root)?;
-    }
-
-    // 5. Apply overlays (common)
-    for overlay_block in &spec.overlays {
-        overlays::apply_overlays(&overlay_block.actions, &staging_root, files_dir, runner).await?;
-    }
-
-    info!("Phase 1 complete: rootfs assembled");
 
     Ok(Phase1Result {
         staging_root,

crates/forge-engine/src/phase2/mod.rs

@@ -10,18 +10,16 @@ use spec_parser::schema::{Target, TargetKind};
 use tracing::info;
 
 use crate::error::ForgeError;
-use crate::tools::ToolRunner;
 
-/// Execute Phase 2: produce the target artifact from the staged rootfs.
+/// Execute Phase 2 for non-QCOW2 targets (OCI, Artifact).
 ///
-/// After building the artifact, if a `push_to` reference is set on a QCOW2 target,
+/// QCOW2 targets are handled directly by the orchestrator in `lib.rs` via
-/// the QCOW2 file is automatically pushed as an OCI artifact.
+/// the prepare/finalize/cleanup flow.
 pub async fn execute(
     target: &Target,
     staging_root: &Path,
     files_dir: &Path,
     output_dir: &Path,
-    runner: &dyn ToolRunner,
 ) -> Result<(), ForgeError> {
     info!(
         target = %target.name,
@@ -33,16 +31,23 @@ pub async fn execute(
         TargetKind::Oci => {
             oci::build_oci(target, staging_root, output_dir)?;
         }
-        TargetKind::Qcow2 => {
-            qcow2::build_qcow2(target, staging_root, output_dir, runner).await?;
-        }
         TargetKind::Artifact => {
             artifact::build_artifact(target, staging_root, output_dir, files_dir)?;
         }
+        TargetKind::Qcow2 => {
+            unreachable!("QCOW2 targets are handled by the orchestrator via prepare/finalize/cleanup");
+        }
     }
 
-    // Auto-push QCOW2 to OCI registry if push_to is set
+    info!(target = %target.name, "Phase 2 complete");
-    if target.kind == TargetKind::Qcow2 {
+    Ok(())
+}
+
+/// Push a QCOW2 file to an OCI registry if `push_to` is set on the target.
+pub async fn push_qcow2_if_configured(
+    target: &Target,
+    output_dir: &Path,
+) -> Result<(), ForgeError> {
     if let Some(ref push_ref) = target.push_to {
         let qcow2_path = output_dir.join(format!("{}.qcow2", target.name));
         info!(
@@ -75,8 +80,6 @@ pub async fn execute(
                 detail: e.to_string(),
             })?;
     }
-    }
 
-    info!(target = %target.name, "Phase 2 complete");
     Ok(())
 }

crates/forge-engine/src/phase2/qcow2.rs

@@ -5,22 +5,39 @@ use spec_parser::schema::Target;
 use crate::error::ForgeError;
 use crate::tools::ToolRunner;
 
-/// Build a QCOW2 VM image, dispatching to the appropriate filesystem backend.
+/// A prepared QCOW2 disk image, ready for Phase 1 population.
+#[derive(Debug)]
+pub enum PreparedQcow2 {
+    Ext4(super::qcow2_ext4::PreparedExt4),
+    Zfs(super::qcow2_zfs::PreparedZfs),
+}
+
+impl PreparedQcow2 {
+    /// The path where the target filesystem is mounted; Phase 1 populates here.
+    pub fn root_mount(&self) -> &Path {
+        match self {
+            PreparedQcow2::Ext4(p) => p.root_mount(),
+            PreparedQcow2::Zfs(p) => p.root_mount(),
+        }
+    }
+}
+
+/// Phase 2 prepare: create disk image, partition/format, mount target filesystem.
 ///
-/// - `"zfs"` (default): ZFS pool with boot environment
+/// Dispatches to ext4 or ZFS based on `target.filesystem`.
-/// - `"ext4"`: GPT+EFI+ext4 with GRUB bootloader
+pub async fn prepare_qcow2(
-pub async fn build_qcow2(
     target: &Target,
-    staging_root: &Path,
     output_dir: &Path,
     runner: &dyn ToolRunner,
-) -> Result<(), ForgeError> {
+) -> Result<PreparedQcow2, ForgeError> {
     match target.filesystem.as_deref().unwrap_or("zfs") {
         "zfs" => {
-            super::qcow2_zfs::build_qcow2_zfs(target, staging_root, output_dir, runner).await
+            let prepared = super::qcow2_zfs::prepare_zfs(target, output_dir, runner).await?;
+            Ok(PreparedQcow2::Zfs(prepared))
         }
         "ext4" => {
-            super::qcow2_ext4::build_qcow2_ext4(target, staging_root, output_dir, runner).await
+            let prepared = super::qcow2_ext4::prepare_ext4(target, output_dir, runner).await?;
+            Ok(PreparedQcow2::Ext4(prepared))
         }
         other => Err(ForgeError::UnsupportedFilesystem {
             fs_type: other.to_string(),
@@ -29,6 +46,31 @@ pub async fn build_qcow2(
     }
 }
 
+/// Phase 2 finalize: install bootloader, unmount.
+pub async fn finalize_qcow2(
+    prepared: &PreparedQcow2,
+    runner: &dyn ToolRunner,
+) -> Result<(), ForgeError> {
+    match prepared {
+        PreparedQcow2::Ext4(p) => super::qcow2_ext4::finalize_ext4(p, runner).await,
+        PreparedQcow2::Zfs(p) => super::qcow2_zfs::finalize_zfs(p, runner).await,
+    }
+}
+
+/// Cleanup: detach loopback, convert raw→qcow2 (if `convert`), remove raw file.
+///
+/// Always call this, even if earlier phases failed.
+pub async fn cleanup_qcow2(
+    prepared: PreparedQcow2,
+    convert: bool,
+    runner: &dyn ToolRunner,
+) -> Result<(), ForgeError> {
+    match prepared {
+        PreparedQcow2::Ext4(p) => super::qcow2_ext4::cleanup_ext4(p, convert, runner).await,
+        PreparedQcow2::Zfs(p) => super::qcow2_zfs::cleanup_zfs(p, convert, runner).await,
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -53,12 +95,8 @@ mod tests {
         let target = make_target(Some("btrfs"));
         let rt = tokio::runtime::Runtime::new().unwrap();
         let result = rt.block_on(async {
-            // We can't actually run the build, but we can test the dispatcher logic
-            // by checking the error for an unsupported filesystem
             let tmpdir = tempfile::tempdir().unwrap();
-            let staging = tempfile::tempdir().unwrap();
 
-            // Create a mock runner that always succeeds
             use crate::tools::{ToolOutput, ToolRunner};
             use std::future::Future;
             use std::pin::Pin;
@@ -80,7 +118,7 @@ mod tests {
                 }
             }
 
-            build_qcow2(&target, staging.path(), tmpdir.path(), &FailRunner).await
+            prepare_qcow2(&target, tmpdir.path(), &FailRunner).await
         });
 
         assert!(result.is_err());

crates/forge-engine/src/phase2/qcow2_ext4.rs

@@ -1,4 +1,4 @@
-use std::path::Path;
+use std::path::{Path, PathBuf};
 
 use spec_parser::schema::Target;
 use tracing::info;
@@ -6,26 +6,33 @@ use tracing::info;
 use crate::error::ForgeError;
 use crate::tools::ToolRunner;
 
-/// Build a QCOW2 VM image from the staged rootfs using ext4+GPT+GRUB.
+/// State for a prepared ext4 QCOW2 disk image, ready for Phase 1 population.
+#[derive(Debug)]
+pub struct PreparedExt4 {
+    pub raw_path: PathBuf,
+    pub qcow2_path: PathBuf,
+    pub device: String,
+    pub efi_part: String,
+    pub root_part: String,
+    pub mount_dir: tempfile::TempDir,
+}
+
+impl PreparedExt4 {
+    /// The path where the root partition is mounted; Phase 1 populates into this.
+    pub fn root_mount(&self) -> &Path {
+        self.mount_dir.path()
+    }
+}
+
+/// Phase 2 prepare: create raw disk, partition, format, and mount the root partition.
 ///
-/// Pipeline:
+/// Returns a `PreparedExt4` whose `root_mount()` is the mounted ext4 root —
-/// 1. Create raw disk image of specified size
+/// Phase 1 should populate the rootfs directly into that directory.
-/// 2. Attach loopback device + partprobe
+pub async fn prepare_ext4(
-/// 3. Create GPT partition table (EFI + root)
-/// 4. Format partitions (FAT32 for EFI, ext4 for root)
-/// 5. Mount root, copy staging rootfs
-/// 6. Mount EFI at /boot/efi
-/// 7. Bind-mount /dev, /proc, /sys
-/// 8. chroot grub-install
-/// 9. chroot grub-mkconfig
-/// 10. Unmount all, detach loopback
-/// 11. Convert raw -> qcow2
-pub async fn build_qcow2_ext4(
     target: &Target,
-    staging_root: &Path,
     output_dir: &Path,
     runner: &dyn ToolRunner,
-) -> Result<(), ForgeError> {
+) -> Result<PreparedExt4, ForgeError> {
     let disk_size = target
         .disk_size
         .as_deref()
@@ -34,46 +41,54 @@ pub async fn build_qcow2_ext4(
     let raw_path = output_dir.join(format!("{}.raw", target.name));
     let qcow2_path = output_dir.join(format!("{}.qcow2", target.name));
     let raw_str = raw_path.to_str().unwrap();
-    let qcow2_str = qcow2_path.to_str().unwrap();
 
     info!(disk_size, "Step 1: Creating raw disk image");
     crate::tools::qemu_img::create_raw(runner, raw_str, disk_size).await?;
 
     info!("Step 2: Attaching loopback device");
     let device = crate::tools::loopback::attach(runner, raw_str).await?;
-
-    // Re-read partition table after attaching loopback
     let _ = crate::tools::loopback::partprobe(runner, &device).await;
 
-    let result = async {
     info!(device = %device, "Step 3: Creating GPT partition table");
     let (efi_part, root_part) =
         crate::tools::partition::create_gpt_efi_root(runner, &device).await?;
 
-        // Re-read partition table after creating partitions
     crate::tools::loopback::partprobe(runner, &device).await?;
 
     info!("Step 4: Formatting partitions");
     crate::tools::partition::mkfs_fat32(runner, &efi_part).await?;
     crate::tools::partition::mkfs_ext4(runner, &root_part).await?;
 
-        // Create a temporary mountpoint for the root partition
     let mount_dir = tempfile::tempdir().map_err(ForgeError::StagingSetup)?;
     let mount_str = mount_dir.path().to_str().unwrap();
 
-        info!("Step 5: Mounting root partition and copying rootfs");
+    info!("Step 5: Mounting root partition at {}", mount_str);
     crate::tools::partition::mount(runner, &root_part, mount_str).await?;
 
-        // Copy staging rootfs into mounted root
+    Ok(PreparedExt4 {
-        copy_rootfs(staging_root, mount_dir.path(), runner).await?;
+        raw_path,
+        qcow2_path,
+        device,
+        efi_part,
+        root_part,
+        mount_dir,
+    })
+}
 
-        info!("Step 6: Mounting EFI partition");
+/// Phase 2 finalize: mount EFI, bind-mount pseudofs, install GRUB, unmount everything.
-        let efi_mount = mount_dir.path().join("boot/efi");
+pub async fn finalize_ext4(
+    prepared: &PreparedExt4,
+    runner: &dyn ToolRunner,
+) -> Result<(), ForgeError> {
+    let mount_str = prepared.mount_dir.path().to_str().unwrap();
+
+    info!("Finalize step 1: Mounting EFI partition");
+    let efi_mount = prepared.mount_dir.path().join("boot/efi");
     std::fs::create_dir_all(&efi_mount)?;
     let efi_mount_str = efi_mount.to_str().unwrap();
-        crate::tools::partition::mount(runner, &efi_part, efi_mount_str).await?;
+    crate::tools::partition::mount(runner, &prepared.efi_part, efi_mount_str).await?;
 
-        info!("Step 7: Bind-mounting /dev, /proc, /sys");
+    info!("Finalize step 2: Bind-mounting /dev, /proc, /sys");
     let dev_mount = format!("{mount_str}/dev");
     let proc_mount = format!("{mount_str}/proc");
     let sys_mount = format!("{mount_str}/sys");
@@ -84,7 +99,7 @@ pub async fn build_qcow2_ext4(
     crate::tools::partition::bind_mount(runner, "/proc", &proc_mount).await?;
     crate::tools::partition::bind_mount(runner, "/sys", &sys_mount).await?;
 
-        info!("Step 8: Installing GRUB bootloader");
+    info!("Finalize step 3: Installing GRUB bootloader");
     runner
         .run(
             "chroot",
@@ -98,7 +113,7 @@ pub async fn build_qcow2_ext4(
         )
         .await?;
 
-        info!("Step 9: Generating GRUB config");
+    info!("Finalize step 4: Generating GRUB config");
     runner
         .run(
             "chroot",
@@ -106,55 +121,43 @@ pub async fn build_qcow2_ext4(
         )
         .await?;
 
-        info!("Step 10: Unmounting");
+    info!("Finalize step 5: Unmounting");
-        // Unmount in reverse order: bind mounts, EFI, root
     crate::tools::partition::umount(runner, &sys_mount).await?;
     crate::tools::partition::umount(runner, &proc_mount).await?;
     crate::tools::partition::umount(runner, &dev_mount).await?;
     crate::tools::partition::umount(runner, efi_mount_str).await?;
     crate::tools::partition::umount(runner, mount_str).await?;
 
-        Ok::<(), ForgeError>(())
-    }
-    .await;
-
-    // Always try to detach loopback, even on error
-    info!("Detaching loopback device");
-    let detach_result = crate::tools::loopback::detach(runner, &device).await;
-
-    result?;
-    detach_result?;
-
-    info!("Step 11: Converting raw -> qcow2");
-    crate::tools::qemu_img::convert_to_qcow2(runner, raw_str, qcow2_str).await?;
-
-    // Clean up raw file
-    std::fs::remove_file(&raw_path).ok();
-
-    info!(path = %qcow2_path.display(), "QCOW2 (ext4) image created");
     Ok(())
 }
 
-/// Copy the staging rootfs into the mounted root partition.
+/// Cleanup: detach loopback, convert raw→qcow2 (if `convert` is true), remove raw file.
 ///
-/// Uses `cp -a` (archive mode) to properly preserve symlinks, permissions,
+/// Always runs, even if earlier phases failed — the loopback device must be detached.
-/// ownership, timestamps, and special files. This is critical for modern
+pub async fn cleanup_ext4(
-/// distros with merged /usr where /lib, /bin, /sbin are symlinks.
+    prepared: PreparedExt4,
-async fn copy_rootfs(
+    convert: bool,
-    src: &Path,
-    dest: &Path,
     runner: &dyn ToolRunner,
 ) -> Result<(), ForgeError> {
-    let src_str = format!("{}/.", src.display());
+    info!("Cleanup: detaching loopback device");
-    let dest_str = dest.to_str().unwrap();
+    let detach_result = crate::tools::loopback::detach(runner, &prepared.device).await;
 
-    runner
+    if convert {
-        .run("cp", &["-a", &src_str, dest_str])
+        let raw_str = prepared.raw_path.to_str().unwrap();
-        .await
+        let qcow2_str = prepared.qcow2_path.to_str().unwrap();
-        .map_err(|_| ForgeError::Qcow2Build {
+
-            step: "copy_rootfs".to_string(),
+        info!("Cleanup: converting raw -> qcow2");
-            detail: format!("cp -a {}/. -> {}", src.display(), dest.display()),
+        crate::tools::qemu_img::convert_to_qcow2(runner, raw_str, qcow2_str).await?;
-        })?;
+    }
+
+    // Clean up raw file
+    std::fs::remove_file(&prepared.raw_path).ok();
+
+    detach_result?;
+
+    if convert {
+        info!(path = %prepared.qcow2_path.display(), "QCOW2 (ext4) image created");
+    }
 
     Ok(())
 }

crates/forge-engine/src/phase2/qcow2_zfs.rs

@@ -1,4 +1,4 @@
-use std::path::Path;
+use std::path::{Path, PathBuf};
 
 use spec_parser::schema::Target;
 use tracing::info;
@@ -6,35 +6,45 @@ use tracing::info;
 use crate::error::ForgeError;
 use crate::tools::ToolRunner;
 
-/// Build a QCOW2 VM image from the staged rootfs.
+/// State for a prepared ZFS QCOW2 disk image, ready for Phase 1 population.
+#[derive(Debug)]
+pub struct PreparedZfs {
+    pub raw_path: PathBuf,
+    pub qcow2_path: PathBuf,
+    pub device: String,
+    pub pool_name: String,
+    pub be_dataset: String,
+    pub bootloader_type: String,
+    pub mount_dir: tempfile::TempDir,
+}
+
+impl PreparedZfs {
+    /// The path where the ZFS boot-environment dataset is mounted;
+    /// Phase 1 populates into this.
+    pub fn root_mount(&self) -> &Path {
+        self.mount_dir.path()
+    }
+}
+
+/// Phase 2 prepare: create raw disk, attach loopback, create ZFS pool + BE, mount.
 ///
-/// Pipeline:
+/// Returns a `PreparedZfs` whose `root_mount()` is the mounted BE dataset —
-/// 1. Create raw disk image of specified size
+/// Phase 1 should populate the rootfs directly into that directory.
-/// 2. Attach loopback device
+pub async fn prepare_zfs(
-/// 3. Create ZFS pool with spec properties
-/// 4. Create boot environment structure (rpool/ROOT/be-1)
-/// 5. Copy staging rootfs into mounted BE
-/// 6. Install bootloader via chroot
-/// 7. Set bootfs property
-/// 8. Export pool, detach loopback
-/// 9. Convert raw -> qcow2
-pub async fn build_qcow2_zfs(
     target: &Target,
-    staging_root: &Path,
     output_dir: &Path,
     runner: &dyn ToolRunner,
-) -> Result<(), ForgeError> {
+) -> Result<PreparedZfs, ForgeError> {
     let disk_size = target
         .disk_size
         .as_deref()
         .ok_or(ForgeError::MissingDiskSize)?;
 
-    let bootloader_type = target.bootloader.as_deref().unwrap_or("uefi");
+    let bootloader_type = target.bootloader.as_deref().unwrap_or("uefi").to_string();
 
     let raw_path = output_dir.join(format!("{}.raw", target.name));
     let qcow2_path = output_dir.join(format!("{}.qcow2", target.name));
     let raw_str = raw_path.to_str().unwrap();
-    let qcow2_str = qcow2_path.to_str().unwrap();
 
     // Collect pool properties
     let pool_props: Vec<(&str, &str)> = target
@@ -48,7 +58,7 @@ pub async fn build_qcow2_zfs(
         })
         .unwrap_or_default();
 
-    let pool_name = "rpool";
+    let pool_name = "rpool".to_string();
     let be_dataset = format!("{pool_name}/ROOT/be-1");
 
     info!(disk_size, "Step 1: Creating raw disk image");
@@ -57,10 +67,8 @@ pub async fn build_qcow2_zfs(
     info!("Step 2: Attaching loopback device");
     let device = crate::tools::loopback::attach(runner, raw_str).await?;
 
-    // Wrap the rest in a closure-like structure so we can clean up on error
-    let result = async {
     info!(device = %device, "Step 3: Creating ZFS pool");
-        crate::tools::zpool::create(runner, pool_name, &device, &pool_props).await?;
+    crate::tools::zpool::create(runner, &pool_name, &device, &pool_props).await?;
 
     info!("Step 4: Creating boot environment structure");
     crate::tools::zfs::create(
@@ -70,80 +78,84 @@ pub async fn build_qcow2_zfs(
     )
     .await?;
 
-        let staging_str = staging_root.to_str().unwrap_or(".");
+    // Mount the BE dataset at a fresh tempdir — not the Phase 1 staging root.
+    let mount_dir = tempfile::tempdir().map_err(ForgeError::StagingSetup)?;
+    let mount_str = mount_dir.path().to_str().unwrap();
+
     crate::tools::zfs::create(
         runner,
         &be_dataset,
-            &[("canmount", "noauto"), ("mountpoint", staging_str)],
+        &[("canmount", "noauto"), ("mountpoint", mount_str)],
     )
     .await?;
 
     crate::tools::zfs::mount(runner, &be_dataset).await?;
 
-        info!("Step 5: Copying staging rootfs into boot environment");
+    info!("Step 5: ZFS BE mounted at {}", mount_str);
-        copy_rootfs(staging_root, staging_root)?;
 
-        info!("Step 6: Installing bootloader");
+    Ok(PreparedZfs {
-        crate::tools::bootloader::install(runner, staging_str, pool_name, bootloader_type).await?;
+        raw_path,
-
+        qcow2_path,
-        info!("Step 7: Setting bootfs property");
+        device,
-        crate::tools::zpool::set(runner, pool_name, "bootfs", &be_dataset).await?;
+        pool_name,
-
+        be_dataset,
-        info!("Step 8: Exporting ZFS pool");
+        bootloader_type,
-        crate::tools::zfs::unmount(runner, &be_dataset).await?;
+        mount_dir,
-        crate::tools::zpool::export(runner, pool_name).await?;
+    })
-
-        Ok::<(), ForgeError>(())
 }
-    .await;
 
-    // Always try to detach loopback, even on error
+/// Phase 2 finalize: install bootloader, set bootfs, unmount + export pool.
-    info!("Detaching loopback device");
+pub async fn finalize_zfs(
-    let detach_result = crate::tools::loopback::detach(runner, &device).await;
+    prepared: &PreparedZfs,
+    runner: &dyn ToolRunner,
+) -> Result<(), ForgeError> {
+    let mount_str = prepared.mount_dir.path().to_str().unwrap();
 
-    // Return the original error if there was one
+    info!("Finalize step 1: Installing bootloader");
-    result?;
+    crate::tools::bootloader::install(
-    detach_result?;
+        runner,
+        mount_str,
+        &prepared.pool_name,
+        &prepared.bootloader_type,
+    )
+    .await?;
 
-    info!("Step 9: Converting raw -> qcow2");
+    info!("Finalize step 2: Setting bootfs property");
+    crate::tools::zpool::set(runner, &prepared.pool_name, "bootfs", &prepared.be_dataset).await?;
+
+    info!("Finalize step 3: Unmounting and exporting ZFS pool");
+    crate::tools::zfs::unmount(runner, &prepared.be_dataset).await?;
+    crate::tools::zpool::export(runner, &prepared.pool_name).await?;
+
+    Ok(())
+}
+
+/// Cleanup: detach loopback, convert raw→qcow2 (if `convert` is true), remove raw file.
+///
+/// Always runs, even if earlier phases failed — the loopback device must be detached.
+pub async fn cleanup_zfs(
+    prepared: PreparedZfs,
+    convert: bool,
+    runner: &dyn ToolRunner,
+) -> Result<(), ForgeError> {
+    info!("Cleanup: detaching loopback device");
+    let detach_result = crate::tools::loopback::detach(runner, &prepared.device).await;
+
+    if convert {
+        let raw_str = prepared.raw_path.to_str().unwrap();
+        let qcow2_str = prepared.qcow2_path.to_str().unwrap();
+
+        info!("Cleanup: converting raw -> qcow2");
         crate::tools::qemu_img::convert_to_qcow2(runner, raw_str, qcow2_str).await?;
+    }
 
     // Clean up raw file
-    std::fs::remove_file(&raw_path).ok();
+    std::fs::remove_file(&prepared.raw_path).ok();
 
-    info!(path = %qcow2_path.display(), "QCOW2 image created");
+    detach_result?;
-    Ok(())
-}
 
-/// Copy the staging rootfs into the mounted BE.
+    if convert {
-/// Since the BE is mounted at the staging root mountpoint, we use a recursive
+        info!(path = %prepared.qcow2_path.display(), "QCOW2 (ZFS) image created");
-/// copy approach for files that need relocation.
-fn copy_rootfs(src: &Path, dest: &Path) -> Result<(), ForgeError> {
-    // In the actual build, the ZFS dataset is mounted at the staging_root path,
-    // so the files are already in place after package installation. This function
-    // handles the case where we need to copy from a temp staging dir into the
-    // mounted ZFS dataset.
-    if src == dest {
-        return Ok(());
-    }
-
-    for entry in walkdir::WalkDir::new(src).follow_links(false) {
-        let entry = entry.map_err(|e| ForgeError::Qcow2Build {
-            step: "copy_rootfs".to_string(),
-            detail: e.to_string(),
-        })?;
-
-        let rel = entry.path().strip_prefix(src).unwrap_or(entry.path());
-        let target = dest.join(rel);
-
-        if entry.path().is_dir() {
-            std::fs::create_dir_all(&target)?;
-        } else if entry.path().is_file() {
-            if let Some(parent) = target.parent() {
-                std::fs::create_dir_all(parent)?;
-            }
-            std::fs::copy(entry.path(), &target)?;
-        }
     }
 
     Ok(())