feat: Defer cleanup of compressed driver source folders

Implements a deferred cleanup mechanism for driver source folders when they are compressed to a WIM and also used for WinPE.

When drivers are compressed, the original source folders are now preserved if they are also needed for WinPE driver injection. A marker file is created in these preserved folders.

A new cleanup step is added after the WinPE media creation to remove these preserved folders, ensuring they are available when needed but not left behind permanently.

FFUDevelopment/BuildFFUVM.ps1

@@ -2728,7 +2728,8 @@ function Copy-Drivers {
                 foreach ($sourceDiskFile in $sourceDiskFiles.Keys) {
                     if (!$sourceDiskFiles[$sourceDiskFile].Contains(",")) {
                         Copy-Item -Path "$infPath\$sourceDiskFile" -Destination $targetPath -Force
-                    } else {
+                    }
+                    else {
                         $subdir = ($sourceDiskFiles[$sourceDiskFile] -split ",")[1]
                         [void](New-Item -Path "$targetPath\$subdir" -ItemType Directory -Force)
                         Copy-Item -Path "$infPath\$subdir\$sourceDiskFile" -Destination "$targetPath\$subdir" -Force
@@ -2740,7 +2741,8 @@ function Copy-Drivers {
                 foreach ($sourceDiskFile in $sourceDiskFiles.Keys) {
                     if (!$sourceDiskFiles[$sourceDiskFile].Contains(",")) {
                         Copy-Item -Path "$infPath\$sourceDiskFile" -Destination $targetPath -Force
-                    } else {
+                    }
+                    else {
                         $subdir = ($sourceDiskFiles[$sourceDiskFile] -split ",")[1]
                         [void](New-Item -Path "$targetPath\$subdir" -ItemType Directory -Force)
                         Copy-Item -Path "$infPath\$subdir\$sourceDiskFile" -Destination "$targetPath\$subdir" -Force
@@ -2902,6 +2904,41 @@ function New-PEMedia {
     WriteLog "Cleaning up $WinPEFFUPath"
     Remove-Item -Path "$WinPEFFUPath" -Recurse -Force
     WriteLog 'Cleanup complete'
+    # Deferred cleanup of preserved driver model folders (only after WinPE Deploy media is created)
+    if ($UseDriversAsPEDrivers -and $CompressDownloadedDriversToWim -and $Deploy -and $CopyPEDrivers) {
+        WriteLog "Beginning deferred cleanup of preserved driver model folders (UseDriversAsPEDrivers + compression scenario)."
+        $removedCount = 0
+        $skippedCount = 0
+        if (Test-Path -Path $DriversFolder) {
+            Get-ChildItem -Path $DriversFolder -Directory -ErrorAction SilentlyContinue | ForEach-Object {
+                $makeDir = $_.FullName
+                Get-ChildItem -Path $makeDir -Directory -ErrorAction SilentlyContinue | ForEach-Object {
+                    $modelDir = $_.FullName
+                    $markerFile = Join-Path -Path $modelDir -ChildPath '__PreservedForPEDrivers.txt'
+                    $leaf = Split-Path -Path $modelDir -Leaf
+                    $wimPath = Join-Path -Path $makeDir -ChildPath ($leaf + '.wim')
+                    if ((Test-Path -Path $markerFile -PathType Leaf) -and (Test-Path -Path $wimPath -PathType Leaf)) {
+                        try {
+                            WriteLog "Removing preserved driver folder: $modelDir (WIM located at $wimPath)"
+                            Remove-Item -Path $modelDir -Recurse -Force -ErrorAction Stop
+                            $removedCount++
+                        }
+                        catch {
+                            WriteLog "Warning: Failed to remove preserved folder $modelDir : $($_.Exception.Message)"
+                            $skippedCount++
+                        }
+                    }
+                    else {
+                        $skippedCount++
+                    }
+                }
+            }
+            WriteLog "Deferred driver cleanup complete. Removed: $removedCount; Skipped: $skippedCount"
+        }
+        else {
+            WriteLog "Drivers folder $DriversFolder not found during deferred cleanup."
+        }
+    }
 }
 
 function Optimize-FFUCaptureDrive {
@@ -4812,14 +4849,16 @@ if ($driversJsonPath -and (Test-Path $driversJsonPath) -and ($InstallDrivers -or
     else {
         WriteLog "Found $($driversToProcess.Count) driver entries to process from $driversJsonPath."
 
+        $preserveSourceOnCompress = ($UseDriversAsPEDrivers -and $CompressDownloadedDriversToWim)
         $taskArguments = @{
-            DriversFolder  = $DriversFolder
-            WindowsRelease = $WindowsRelease
-            WindowsArch    = $WindowsArch
-            WindowsVersion = $WindowsVersion 
-            Headers        = $Headers
-            UserAgent      = $UserAgent
-            CompressToWim  = $CompressDownloadedDriversToWim
+            DriversFolder            = $DriversFolder
+            WindowsRelease           = $WindowsRelease
+            WindowsArch              = $WindowsArch
+            WindowsVersion           = $WindowsVersion 
+            Headers                  = $Headers
+            UserAgent                = $UserAgent
+            CompressToWim            = $CompressDownloadedDriversToWim
+            PreserveSourceOnCompress = $preserveSourceOnCompress
         }
         
         WriteLog "Starting parallel driver processing using Invoke-ParallelProcessing..."

FFUDevelopment/FFU.Common/FFU.Common.Drivers.psm1

@@ -22,7 +22,10 @@ function Compress-DriverFolderToWim {
         [string]$WimName, # Optional, defaults to folder name
 
         [Parameter()]
-        [string]$WimDescription # Optional, defaults to folder name
+        [string]$WimDescription, # Optional, defaults to folder name
+
+        [Parameter()]
+        [bool]$PreserveSource = $false # When $true, do not delete source folder; create marker for deferred cleanup
     )
 
     WriteLog "Starting compression of folder '$SourceFolderPath' to '$DestinationWimPath'."
@@ -66,14 +69,29 @@ function Compress-DriverFolderToWim {
             WriteLog "Successfully compressed '$SourceFolderPath' to '$DestinationWimPath' using dism.exe."
             
             # Remove the source folder after successful compression
-            WriteLog "Removing source driver folder: $SourceFolderPath"
-            try {
-                Remove-Item -Path $SourceFolderPath -Recurse -Force -ErrorAction Stop
-                WriteLog "Successfully removed source folder '$SourceFolderPath'."
+            if ($PreserveSource) {
+                WriteLog "Preserving source driver folder for deferred WinPE driver harvesting: $SourceFolderPath"
+                try {
+                    $markerFile = Join-Path -Path $SourceFolderPath -ChildPath '__PreservedForPEDrivers.txt'
+                    if (-not (Test-Path -Path $markerFile -PathType Leaf)) {
+                        New-Item -Path $markerFile -ItemType File -Force | Out-Null
+                        WriteLog "Created preservation marker file: $markerFile"
+                    }
+                }
+                catch {
+                    WriteLog "Warning: Failed to create preservation marker in $SourceFolderPath. Error: $($_.Exception.Message)"
+                }
             }
-            catch {
-                WriteLog "Warning: Failed to remove source folder '$SourceFolderPath'. Error: $($_.Exception.Message)"
-                # Do not fail the whole operation, just log a warning.
+            else {
+                WriteLog "Removing source driver folder: $SourceFolderPath"
+                try {
+                    Remove-Item -Path $SourceFolderPath -Recurse -Force -ErrorAction Stop
+                    WriteLog "Successfully removed source folder '$SourceFolderPath'."
+                }
+                catch {
+                    WriteLog "Warning: Failed to remove source folder '$SourceFolderPath'. Error: $($_.Exception.Message)"
+                    # Do not fail the whole operation, just log a warning.
+                }
             }
 
             return $true # Indicate success

FFUDevelopment/FFU.Common/FFU.Common.Parallel.psm1

@@ -209,7 +209,8 @@ function Invoke-ParallelProcessing {
                                     -Headers $localJobArgs['Headers'] `
                                     -UserAgent $localJobArgs['UserAgent'] `
                                     -ProgressQueue $localProgressQueue `
-                                    -CompressToWim $localJobArgs['CompressToWim']
+                                    -CompressToWim $localJobArgs['CompressToWim'] `
+                                    -PreserveSourceOnCompress $localJobArgs['PreserveSourceOnCompress']
                             }
                             'Dell' {
                                 $taskResult = Save-DellDriversTask -DriverItemData $currentItem `
@@ -217,7 +218,8 @@ function Invoke-ParallelProcessing {
                                     -WindowsArch $localJobArgs['WindowsArch'] `
                                     -WindowsRelease $localJobArgs['WindowsRelease'] `
                                     -ProgressQueue $localProgressQueue `
-                                    -CompressToWim $localJobArgs['CompressToWim']
+                                    -CompressToWim $localJobArgs['CompressToWim'] `
+                                    -PreserveSourceOnCompress $localJobArgs['PreserveSourceOnCompress']
                             }
                             'HP' {
                                 $taskResult = Save-HPDriversTask -DriverItemData $currentItem `
@@ -226,7 +228,8 @@ function Invoke-ParallelProcessing {
                                     -WindowsRelease $localJobArgs['WindowsRelease'] `
                                     -WindowsVersion $localJobArgs['WindowsVersion'] `
                                     -ProgressQueue $localProgressQueue `
-                                    -CompressToWim $localJobArgs['CompressToWim']
+                                    -CompressToWim $localJobArgs['CompressToWim'] `
+                                    -PreserveSourceOnCompress $localJobArgs['PreserveSourceOnCompress']
                             }
                             'Lenovo' {
                                 $taskResult = Save-LenovoDriversTask -DriverItemData $currentItem `
@@ -235,7 +238,8 @@ function Invoke-ParallelProcessing {
                                     -Headers $localJobArgs['Headers'] `
                                     -UserAgent $localJobArgs['UserAgent'] `
                                     -ProgressQueue $localProgressQueue `
-                                    -CompressToWim $localJobArgs['CompressToWim']
+                                    -CompressToWim $localJobArgs['CompressToWim'] `
+                                    -PreserveSourceOnCompress $localJobArgs['PreserveSourceOnCompress']
                             }
                             default {
                                 $unsupportedMakeMessage = "Error: Unsupported Make '$make' for driver download."

FFUDevelopment/FFUUI.Core/FFUUI.Core.Drivers.Dell.psm1

@@ -167,10 +167,12 @@ function Save-DellDriversTask {
         [string]$WindowsArch,
         [Parameter(Mandatory = $true)]
         [int]$WindowsRelease,
-        [Parameter()] # Made optional
+        [Parameter()]
         [System.Collections.Concurrent.ConcurrentQueue[hashtable]]$ProgressQueue = $null, # Default to null
         [Parameter()]
-        [bool]$CompressToWim = $false # New parameter for compression
+        [bool]$CompressToWim = $false, # New parameter for compression
+        [Parameter()]
+        [bool]$PreserveSourceOnCompress = $false
     )
         
     $modelName = $DriverItemData.Model
@@ -201,7 +203,7 @@ function Save-DellDriversTask {
                 WriteLog "Attempting compression of existing folder '$sourceFolderPath' to '$wimFilePath'."
                 if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $modelName -Status "Compressing existing..." }
                 try {
-                    Compress-DriverFolderToWim -SourceFolderPath $sourceFolderPath -DestinationWimPath $wimFilePath -WimName $modelName -WimDescription "Drivers for $modelName" -ErrorAction Stop
+                    Compress-DriverFolderToWim -SourceFolderPath $sourceFolderPath -DestinationWimPath $wimFilePath -WimName $modelName -WimDescription "Drivers for $modelName" -PreserveSource:$PreserveSourceOnCompress -ErrorAction Stop
                     $existingDriver.Status = "Already downloaded & Compressed"
                     $existingDriver.DriverPath = Join-Path -Path $make -ChildPath "$($modelName).wim"
                     $existingDriver.Success = $true
@@ -664,7 +666,7 @@ function Save-DellDriversTask {
             $driverRelativePath = Join-Path -Path $make -ChildPath $wimFileName # Update relative path to the WIM file
             WriteLog "Compressing '$modelPath' to '$destinationWimPath'..."
             try {
-                $compressResult = Compress-DriverFolderToWim -SourceFolderPath $modelPath -DestinationWimPath $destinationWimPath -WimName $modelName -WimDescription $modelName -ErrorAction Stop
+                $compressResult = Compress-DriverFolderToWim -SourceFolderPath $modelPath -DestinationWimPath $destinationWimPath -WimName $modelName -WimDescription $modelName -PreserveSource:$PreserveSourceOnCompress -ErrorAction Stop
                 if ($compressResult) {
                     WriteLog "Compression successful for '$modelName'."
                     $status = "Completed & Compressed"

FFUDevelopment/FFUUI.Core/FFUUI.Core.Drivers.HP.psm1

@@ -118,7 +118,9 @@ function Save-HPDriversTask {
         [Parameter()] # Made optional
         [System.Collections.Concurrent.ConcurrentQueue[hashtable]]$ProgressQueue = $null, # Default to null
         [Parameter()]
-        [bool]$CompressToWim = $false # New parameter for compression
+        [bool]$CompressToWim = $false, # New parameter for compression
+        [Parameter()]
+        [bool]$PreserveSourceOnCompress = $false
     )
             
     $modelName = $DriverItemData.Model
@@ -150,7 +152,7 @@ function Save-HPDriversTask {
                 WriteLog "Attempting compression of existing folder '$sourceFolderPath' to '$wimFilePath'."
                 if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $identifier -Status "Compressing existing..." }
                 try {
-                    Compress-DriverFolderToWim -SourceFolderPath $sourceFolderPath -DestinationWimPath $wimFilePath -WimName $identifier -WimDescription "Drivers for $identifier" -ErrorAction Stop
+                    Compress-DriverFolderToWim -SourceFolderPath $sourceFolderPath -DestinationWimPath $wimFilePath -WimName $identifier -WimDescription "Drivers for $identifier" -PreserveSource:$PreserveSourceOnCompress -ErrorAction Stop
                     $existingDriver.Status = "Already downloaded & Compressed"
                     $existingDriver.DriverPath = Join-Path -Path $make -ChildPath "$($sanitizedModelName).wim"
                     $existingDriver.Success = $true
@@ -362,7 +364,7 @@ function Save-HPDriversTask {
             $wimFilePath = Join-Path -Path $hpDriversBaseFolder -ChildPath "$($identifier).wim"
             WriteLog "Compressing '$modelSpecificFolder' to '$wimFilePath'..."
             try {
-                Compress-DriverFolderToWim -SourceFolderPath $modelSpecificFolder -DestinationWimPath $wimFilePath -WimName $identifier -WimDescription "Drivers for $identifier" -ErrorAction Stop
+                Compress-DriverFolderToWim -SourceFolderPath $modelSpecificFolder -DestinationWimPath $wimFilePath -WimName $identifier -WimDescription "Drivers for $identifier" -PreserveSource:$PreserveSourceOnCompress -ErrorAction Stop
                 WriteLog "Compression successful for '$identifier'."
                 $finalStatus = "Completed & Compressed"
                 $driverRelativePath = Join-Path -Path $make -ChildPath "$($identifier).wim" # Update relative path to the WIM

FFUDevelopment/FFUUI.Core/FFUUI.Core.Drivers.Lenovo.psm1

@@ -94,10 +94,12 @@ function Save-LenovoDriversTask {
         [hashtable]$Headers,
         [Parameter(Mandatory = $true)]
         [string]$UserAgent,
-        [Parameter()] # Made optional
+        [Parameter()]
         [System.Collections.Concurrent.ConcurrentQueue[hashtable]]$ProgressQueue = $null,
         [Parameter()]
-        [bool]$CompressToWim = $false
+        [bool]$CompressToWim = $false,
+        [Parameter()]
+        [bool]$PreserveSourceOnCompress = $false
     )
             
     # The Model property from the UI already contains the combined "ProductName (MachineType)" string
@@ -133,7 +135,7 @@ function Save-LenovoDriversTask {
                 WriteLog "Attempting compression of existing folder '$sourceFolderPath' to '$wimFilePath'."
                 if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $identifier -Status "Compressing existing..." }
                 try {
-                    Compress-DriverFolderToWim -SourceFolderPath $sourceFolderPath -DestinationWimPath $wimFilePath -WimName $identifier -WimDescription "Drivers for $identifier" -ErrorAction Stop
+                    Compress-DriverFolderToWim -SourceFolderPath $sourceFolderPath -DestinationWimPath $wimFilePath -WimName $identifier -WimDescription "Drivers for $identifier" -PreserveSource:$PreserveSourceOnCompress -ErrorAction Stop
                     $existingDriver.Status = "Already downloaded & Compressed"
                     $existingDriver.DriverPath = Join-Path -Path $make -ChildPath "$($sanitizedIdentifier).wim"
                     $existingDriver.Success = $true
@@ -424,7 +426,7 @@ function Save-LenovoDriversTask {
             $driverRelativePath = Join-Path -Path $make -ChildPath $wimFileName # Update relative path to the WIM file
             WriteLog "Compressing '$modelPath' to '$destinationWimPath'..."
             try {
-                $compressResult = Compress-DriverFolderToWim -SourceFolderPath $modelPath -DestinationWimPath $destinationWimPath -WimName $identifier -WimDescription $identifier -ErrorAction Stop
+                $compressResult = Compress-DriverFolderToWim -SourceFolderPath $modelPath -DestinationWimPath $destinationWimPath -WimName $identifier -WimDescription $identifier -PreserveSource:$PreserveSourceOnCompress -ErrorAction Stop
                 if ($compressResult) {
                     WriteLog "Compression successful for '$identifier'."
                     $status = "Completed & Compressed"

FFUDevelopment/FFUUI.Core/FFUUI.Core.Drivers.Microsoft.psm1

@@ -94,7 +94,9 @@ function Save-MicrosoftDriversTask {
         [Parameter()] # Made optional
         [System.Collections.Concurrent.ConcurrentQueue[hashtable]]$ProgressQueue = $null, # Default to null
         [Parameter()]
-        [bool]$CompressToWim = $false # New parameter for compression
+        [bool]$CompressToWim = $false, # New parameter for compression
+        [Parameter()]
+        [bool]$PreserveSourceOnCompress = $false
         # REMOVED: UI-related parameters
     )
         
@@ -125,7 +127,7 @@ function Save-MicrosoftDriversTask {
                 WriteLog "Attempting compression of existing folder '$sourceFolderPath' to '$wimFilePath'."
                 if ($null -ne $ProgressQueue) { Invoke-ProgressUpdate -ProgressQueue $ProgressQueue -Identifier $modelName -Status "Compressing existing..." }
                 try {
-                    Compress-DriverFolderToWim -SourceFolderPath $sourceFolderPath -DestinationWimPath $wimFilePath -WimName $modelName -WimDescription "Drivers for $modelName" -ErrorAction Stop
+                    Compress-DriverFolderToWim -SourceFolderPath $sourceFolderPath -DestinationWimPath $wimFilePath -WimName $modelName -WimDescription "Drivers for $modelName" -PreserveSource:$PreserveSourceOnCompress -ErrorAction Stop
                     $existingDriver.Status = "Already downloaded & Compressed"
                     $existingDriver.DriverPath = Join-Path -Path $make -ChildPath "$($modelName).wim"
                     $existingDriver.Success = $true
@@ -377,7 +379,7 @@ function Save-MicrosoftDriversTask {
                     WriteLog "Compressing '$modelPath' to '$destinationWimPath'..."
                     try {
                         # Use the function from the imported common module
-                        $compressResult = Compress-DriverFolderToWim -SourceFolderPath $modelPath -DestinationWimPath $destinationWimPath -WimName $modelName -WimDescription $modelName -ErrorAction Stop
+                        $compressResult = Compress-DriverFolderToWim -SourceFolderPath $modelPath -DestinationWimPath $destinationWimPath -WimName $modelName -WimDescription $modelName -PreserveSource:$PreserveSourceOnCompress -ErrorAction Stop
                         if ($compressResult) {
                             WriteLog "Compression successful for '$modelName'."
                             $status = "Completed & Compressed"

FFUDevelopment/FFUUI.Core/FFUUI.Core.Drivers.psm1

@@ -571,6 +571,8 @@ function Invoke-DownloadSelectedDrivers {
     $localHeaders = $coreStaticVars.Headers
     $localUserAgent = $coreStaticVars.UserAgent
     $compressDrivers = $State.Controls.chkCompressDriversToWIM.IsChecked
+    # Determine if we must preserve source folders (used later for PE driver harvesting)
+    $preserveSource = ($State.Controls.chkUseDriversAsPEDrivers.IsChecked -and $State.Controls.chkCompressDriversToWIM.IsChecked)
 
     $State.Controls.txtStatus.Text = "Processing all selected drivers..."
     WriteLog "Processing all selected drivers: $($selectedDrivers.Model -join ', ')"
@@ -620,15 +622,16 @@ function Invoke-DownloadSelectedDrivers {
             return
         }
     }
-
+    $preserveSource = ($State.Controls.chkUseDriversAsPEDrivers.IsChecked -and $State.Controls.chkCompressDriversToWIM.IsChecked)
     $taskArguments = @{
-        DriversFolder  = $localDriversFolder
-        WindowsRelease = $localWindowsRelease
-        WindowsArch    = $localWindowsArch
-        WindowsVersion = $localWindowsVersion
-        Headers        = $localHeaders
-        UserAgent      = $localUserAgent
-        CompressToWim  = $compressDrivers
+        DriversFolder            = $localDriversFolder
+        WindowsRelease           = $localWindowsRelease
+        WindowsArch              = $localWindowsArch
+        WindowsVersion           = $localWindowsVersion
+        Headers                  = $localHeaders
+        UserAgent                = $localUserAgent
+        CompressToWim            = $compressDrivers
+        PreserveSourceOnCompress = $preserveSource
     }
 
     $parallelResults = Invoke-ParallelProcessing -ItemsToProcess $selectedDrivers `