feat(template): cleanup-after-update + ruff fixes

Move disk cleanup out of the Windows Update loop:
- PS scripts: extract Invoke-DiskCleanup function, add -CleanupOnly switch
  for standalone post-update cleanup pass.
- Python command: always pass -SkipCleanup:$true during WU loop; after
  loop exits 0, run a separate -CleanupOnly invocation (skipped if
  --skip-cleanup). Fails fast if cleanup-only returns nonzero.
- Tests: 3 new tests for cleanup-runs-after-loop, skip-cleanup suppresses
  it, and cleanup-failure propagates as ClickException.
- ruff --fix: remove unused shutil import, sort imports, fix UP037/UP035.

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

src/ci_orchestrator/commands/template.py

@@ -57,7 +57,6 @@ from __future__ import annotations
 
 import fnmatch
 import json
-import shutil
 import socket
 import subprocess
 import sys
@@ -68,8 +67,8 @@ from pathlib import Path
 
 import click
 
-from ci_orchestrator.backends.workstation import WorkstationVmrunBackend
 from ci_orchestrator.backends.protocol import VmHandle
+from ci_orchestrator.backends.workstation import WorkstationVmrunBackend
 from ci_orchestrator.config import load_config
 from ci_orchestrator.transport.errors import TransportError
 from ci_orchestrator.transport.winrm import WinRmResult, WinRmTransport
@@ -324,7 +323,7 @@ def _run_with_heartbeat(
     def _worker() -> None:
         try:
             result_holder.append(transport.run(script, check=check))
-        except BaseException as exc:  # noqa: BLE001
+        except BaseException as exc:
             exc_holder.append(exc)
         finally:
             stop_event.set()
@@ -507,8 +506,7 @@ def template_prepare_win(
             args_parts.append("-SkipWindowsUpdate:$true")
         if skip_tier2:
             args_parts.append("-SkipTier2:$true")
-        if skip_cleanup:
+        args_parts.append("-SkipCleanup:$true")  # always skip during WU loop; cleanup runs separately after
-            args_parts.append("-SkipCleanup:$true")
         args_str = " ".join(args_parts)
 
         # Run toolchain script — loop handles Windows Update reboot (exit 3010)
@@ -569,6 +567,33 @@ def template_prepare_win(
             _wait_winrm(transport, winrm_timeout)
             click.echo("[prepare-win] WinRM ready after reboot.")
 
+        # Post-update disk cleanup — runs once after all WU reboots complete
+        if not skip_cleanup:
+            click.echo("[prepare-win] Running post-update disk cleanup (-CleanupOnly) ...")
+            cleanup_invoke = (
+                "Set-ExecutionPolicy Bypass -Scope Process -Force\n"
+                "$ec = 0\n"
+                "try {\n"
+                f"    & '{e(guest_script)}' -BuildPassword '{e(build_password)}'"
+                f" -BuildUsername '{e(build_username)}'"
+                f" -AdminPassword '{e(admin_password)}' -CleanupOnly:$true\n"
+                "    $ec = if ($null -eq $LASTEXITCODE) { 0 } else { $LASTEXITCODE }\n"
+                "} catch {\n"
+                "    Write-Error $_\n"
+                "    $ec = 1\n"
+                "}\n"
+                '"CI_EXITCODE:$ec"\n'
+            )
+            cleanup_result = _run_with_heartbeat(transport, cleanup_invoke, check=False)
+            cleanup_ec = _parse_exit_marker(cleanup_result.stdout)
+            if cleanup_result.stdout.strip():
+                click.echo(cleanup_result.stdout.strip())
+            if cleanup_result.stderr.strip():
+                click.echo(cleanup_result.stderr.strip(), err=True)
+            if cleanup_ec != 0:
+                raise click.ClickException(f"Disk cleanup failed (exit {cleanup_ec}).")
+            click.echo("[prepare-win] Post-update cleanup complete.")
+
         # Install CI-StaticIp scheduled task
         if static_ip_path is not None:
             click.echo("[prepare-win] Installing CI-StaticIp scheduled task ...")

src/ci_orchestrator/ip_pool.py

@@ -39,9 +39,9 @@ from __future__ import annotations
 
 import json
 import time
+from collections.abc import Generator
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Generator
 
 
 class IpSlotPool:

template/Install-CIToolchain-WinBuild2022.ps1

@@ -148,6 +148,11 @@ param(
     # Skip disk cleanup step (cleanmgr + DISM + cache clear) — speeds up re-runs during dev/debug
     [switch] $SkipCleanup,
 
+    # Run disk cleanup only — skip all provisioning steps and jump straight to the
+    # cleanup section. Intended for use by Prepare-WinBuild2022.ps1 after the Windows
+    # Update loop completes, so cleanup always runs on the fully-updated system.
+    [switch] $CleanupOnly,
+
     # Administrator password — used only to write DefaultPassword for autologin.
     # Optional: if empty, DefaultPassword registry value is left as-is (Deploy may have set it).
     [string] $AdminPassword = '',
@@ -231,6 +236,85 @@ function Assert-Hash {
     Write-Host "  [OK] SHA256 verified: $Label" -ForegroundColor Green
 }
 
+function Invoke-DiskCleanup {
+    Write-Step "Cleaning up disk"
+
+    $cleanKey = 'HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\Explorer\VolumeCaches'
+    $sageset  = 1
+    $cleanCategories = @(
+        'Active Setup Temp Folders',
+        'BranchCache',
+        'Downloaded Program Files',
+        'Internet Cache Files',
+        'Memory Dump Files',
+        'Old ChkDsk Files',
+        'Previous Installations',
+        'Recycle Bin',
+        'Service Pack Cleanup',
+        'Setup Log Files',
+        'System error memory dump files',
+        'System error minidump files',
+        'Temporary Files',
+        'Temporary Setup Files',
+        'Thumbnail Cache',
+        'Update Cleanup',
+        'Upgrade Discarded Files',
+        'Windows Defender',
+        'Windows Error Reporting Archive Files',
+        'Windows Error Reporting Files',
+        'Windows Error Reporting Queue Files',
+        'Windows Error Reporting Temp Files',
+        'Windows ESD installation files',
+        'Windows Upgrade Log Files'
+    )
+    foreach ($cat in $cleanCategories) {
+        $keyPath = "$cleanKey\$cat"
+        if (Test-Path $keyPath) {
+            Set-ItemProperty -Path $keyPath -Name "StateFlags$('{0:D4}' -f $sageset)" -Value 2 -Type DWord -Force -ErrorAction SilentlyContinue
+        }
+    }
+    Write-Host "Running cleanmgr /sagerun:$sageset (may take a few minutes)..."
+    $proc = Start-Process -FilePath 'cleanmgr.exe' -ArgumentList "/sagerun:$sageset" -Wait -PassThru
+    Write-Host "cleanmgr exited (code $($proc.ExitCode))."
+
+    Write-Host "Stopping WU-adjacent services before cache clear..."
+    foreach ($svc in 'wuauserv', 'UsoSvc', 'bits', 'dosvc', 'TrustedInstaller') {
+        Stop-Service -Name $svc -Force -ErrorAction SilentlyContinue
+    }
+    $sdDownload = 'C:\Windows\SoftwareDistribution\Download'
+    Remove-Item $sdDownload -Recurse -Force -ErrorAction SilentlyContinue
+    New-Item -ItemType Directory -Path $sdDownload -Force | Out-Null
+    Write-Host "Windows Update download cache cleared."
+
+    Remove-Item 'C:\Windows\Logs\CBS\*' -Force -ErrorAction SilentlyContinue
+    Remove-Item 'C:\Windows\Temp\*'     -Recurse -Force -ErrorAction SilentlyContinue
+    Remove-Item "$env:TEMP\*"           -Recurse -Force -ErrorAction SilentlyContinue
+    Remove-Item 'C:\Windows\Prefetch\*' -Force -ErrorAction SilentlyContinue
+
+    Write-Host "Running DISM component store cleanup (StartComponentCleanup)..."
+    $dism = Start-Process -FilePath 'dism.exe' `
+        -ArgumentList '/Online','/Cleanup-Image','/StartComponentCleanup','/ResetBase' `
+        -Wait -PassThru
+    Write-Host "DISM exited (code $($dism.ExitCode))."
+    Write-Host "Disk cleanup complete."
+
+    foreach ($svc in 'wuauserv', 'UsoSvc') {
+        Set-Service -Name $svc -StartupType Disabled -ErrorAction SilentlyContinue
+    }
+    Set-ItemProperty -Path 'HKLM:\SYSTEM\CurrentControlSet\Services\wuauserv' `
+        -Name Start -Value 4 -Type DWord -Force -ErrorAction SilentlyContinue
+
+    $sdFiles = Get-ChildItem 'C:\Windows\SoftwareDistribution\Download' -Recurse -File -Force -ErrorAction SilentlyContinue
+    $sdBytes = [long]0
+    if ($sdFiles) { foreach ($f in $sdFiles) { $sdBytes += $f.Length } }
+    $sdMB = [math]::Round($sdBytes / 1MB, 1)
+    Write-Host "  SoftwareDistribution\Download: $(@($sdFiles).Count) file(s), ${sdMB} MB remaining (WaaSMedicSvc best-effort, WU is disabled)."
+
+    Assert-Step 'Cleanup' 'wuauserv StartType Disabled after cache clear' {
+        (Get-Service wuauserv -ErrorAction Stop).StartType -eq 'Disabled'
+    }
+}
+
 # ── Pinned SHA256 hashes for downloaded installers (§1.3) ─────────────────────
 # Fill in the hash values below. Get them by:
 #   (Invoke-WebRequest '<URL>' -UseBasicParsing -OutFile 'tmp.bin'; (Get-FileHash 'tmp.bin' -Algorithm SHA256).Hash)
@@ -272,6 +356,18 @@ $script:Hashes = @{
     # vcpkg: cloned from GitHub via TLS; pin integrity with $VcpkgRef (commit SHA)
 }
 
+# ── CleanupOnly early-exit ────────────────────────────────────────────────────
+if ($CleanupOnly) {
+    if ($SkipCleanup) {
+        Write-Host "[Setup] -CleanupOnly + -SkipCleanup: nothing to do." -ForegroundColor Yellow
+        exit 0
+    }
+    Write-Host "[Setup] CleanupOnly mode — running disk cleanup on fully-updated system." -ForegroundColor Cyan
+    Invoke-DiskCleanup
+    Write-Host "[Setup] CleanupOnly complete." -ForegroundColor Green
+    exit 0
+}
+
 # ── Pre-flight: remove temp files from any previous partial run ───────────────
 # These files are normally deleted at the end of each step that creates them.
 # A previous interrupted run may have left them behind — remove before starting.
@@ -1527,114 +1623,9 @@ Write-Host "Tier-2 Toolchain installation complete." -ForegroundColor Green
 # ── Cleanup ───────────────────────────────────────────────────────────────────
 if ($SkipCleanup) {
     Write-Host "`n[Setup] Skipping disk cleanup (-SkipCleanup switch set)." -ForegroundColor Yellow
+} else {
+    Invoke-DiskCleanup
 }
-else {
-Write-Step "Cleaning up disk"
-
-# 1. Windows Disk Cleanup (cleanmgr) — run with /sagerun:1 using preset flags
-#    First register the flags via /sageset:1 in the registry (silent, no UI)
-$cleanKey = 'HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\Explorer\VolumeCaches'
-$sageset  = 1
-$cleanCategories = @(
-    'Active Setup Temp Folders',
-    'BranchCache',
-    'Downloaded Program Files',
-    'Internet Cache Files',
-    'Memory Dump Files',
-    'Old ChkDsk Files',
-    'Previous Installations',
-    'Recycle Bin',
-    'Service Pack Cleanup',
-    'Setup Log Files',
-    'System error memory dump files',
-    'System error minidump files',
-    'Temporary Files',
-    'Temporary Setup Files',
-    'Thumbnail Cache',
-    'Update Cleanup',
-    'Upgrade Discarded Files',
-    'Windows Defender',
-    'Windows Error Reporting Archive Files',
-    'Windows Error Reporting Files',
-    'Windows Error Reporting Queue Files',
-    'Windows Error Reporting Temp Files',
-    'Windows ESD installation files',
-    'Windows Upgrade Log Files'
-)
-foreach ($cat in $cleanCategories) {
-    $keyPath = "$cleanKey\$cat"
-    if (Test-Path $keyPath) {
-        Set-ItemProperty -Path $keyPath -Name "StateFlags$('{0:D4}' -f $sageset)" -Value 2 -Type DWord -Force -ErrorAction SilentlyContinue
-    }
-}
-Write-Host "Running cleanmgr /sagerun:$sageset (may take a few minutes)..."
-$proc = Start-Process -FilePath 'cleanmgr.exe' -ArgumentList "/sagerun:$sageset" -Wait -PassThru
-Write-Host "cleanmgr exited (code $($proc.ExitCode))."
-
-# 2. Clear Windows Update cache
-# Stop all services that might hold file locks inside SoftwareDistribution:
-#   wuauserv   — Windows Update (already Disabled from Step 6b, stop is idempotent)
-#   UsoSvc     — Update Orchestrator (already Disabled from Step 6b)
-#   bits       — Background Intelligent Transfer Service (downloads WU payloads)
-#   dosvc      — Delivery Optimization (peer-to-peer WU cache)
-#   TrustedInstaller — Windows Modules Installer (unpacks .cab/.msu update packages)
-Write-Host "Stopping WU-adjacent services before cache clear..."
-foreach ($svc in 'wuauserv', 'UsoSvc', 'bits', 'dosvc', 'TrustedInstaller') {
-    Stop-Service -Name $svc -Force -ErrorAction SilentlyContinue
-}
-
-# Delete and recreate the Download folder — more reliable than 'Remove-Item *'
-# because locked files inside subdirectories silently fail with wildcard delete.
-$sdDownload = 'C:\Windows\SoftwareDistribution\Download'
-Remove-Item $sdDownload -Recurse -Force -ErrorAction SilentlyContinue
-New-Item -ItemType Directory -Path $sdDownload -Force | Out-Null
-Write-Host "Windows Update download cache cleared."
-
-# 3. Clear CBS / component store logs
-Remove-Item 'C:\Windows\Logs\CBS\*' -Force -ErrorAction SilentlyContinue
-
-# 4. Clear TEMP folders
-Remove-Item 'C:\Windows\Temp\*'          -Recurse -Force -ErrorAction SilentlyContinue
-Remove-Item "$env:TEMP\*"                -Recurse -Force -ErrorAction SilentlyContinue
-
-# 5. Clear Prefetch
-Remove-Item 'C:\Windows\Prefetch\*'      -Force -ErrorAction SilentlyContinue
-
-# 6. Run DISM component store cleanup (removes superseded update components)
-Write-Host "Running DISM component store cleanup (StartComponentCleanup)..."
-$dism = Start-Process -FilePath 'dism.exe' `
-    -ArgumentList '/Online','/Cleanup-Image','/StartComponentCleanup','/ResetBase' `
-    -Wait -PassThru
-Write-Host "DISM exited (code $($dism.ExitCode))."
-
-Write-Host "Disk cleanup complete."
-
-# DISM StartComponentCleanup / WaaSMedicSvc may reset service StartType during cleanup.
-# Re-enforce Disabled on both WU services after DISM completes.
-foreach ($svc in 'wuauserv', 'UsoSvc') {
-    Set-Service -Name $svc -StartupType Disabled -ErrorAction SilentlyContinue
-}
-Set-ItemProperty -Path 'HKLM:\SYSTEM\CurrentControlSet\Services\wuauserv' `
-    -Name Start -Value 4 -Type DWord -Force -ErrorAction SilentlyContinue
-
-# Validation — leftover artifacts from this run should be gone
-# Note: do NOT assert SoftwareDistribution\Download empty.
-# WaaSMedicSvc (Windows Update Medic Service) is tamper-protected — cannot be stopped —
-# and actively restores WU files immediately after deletion. Files reappearing here is
-# normal and harmless: WU is permanently disabled by Step 6b (service Disabled + GPO keys).
-# Report remaining size as informational only.
-$sdFiles = Get-ChildItem 'C:\Windows\SoftwareDistribution\Download' -Recurse -File -Force -ErrorAction SilentlyContinue
-# Avoid Measure-Object .Sum: under Set-StrictMode -Version Latest, PS 5.1 treats the
-# nullable Sum property as absent when the pipe is empty, throwing "property not found".
-$sdBytes = [long]0
-if ($sdFiles) { foreach ($f in $sdFiles) { $sdBytes += $f.Length } }
-$sdMB = [math]::Round($sdBytes / 1MB, 1)
-Write-Host "  SoftwareDistribution\Download: $(@($sdFiles).Count) file(s), ${sdMB} MB remaining (WaaSMedicSvc best-effort, WU is disabled)."
-
-Assert-Step 'Cleanup' 'wuauserv StartType Disabled after cache clear' {
-    (Get-Service wuauserv -ErrorAction Stop).StartType -eq 'Disabled'
-}
-} # end if (-not $SkipCleanup)
 
 # ── Pre-Final re-affirmation ──────────────────────────────────────────────────
 # Steps 7-10 (installs, cleanup, DISM) can take hours. During that window background

template/Install-CIToolchain-WinBuild2025.ps1

@@ -148,6 +148,11 @@ param(
     # Skip disk cleanup step (cleanmgr + DISM + cache clear) — speeds up re-runs during dev/debug
     [switch] $SkipCleanup,
 
+    # Run disk cleanup only — skip all provisioning steps and jump straight to the
+    # cleanup section. Intended for use by Prepare-WinBuild2025.ps1 after the Windows
+    # Update loop completes, so cleanup always runs on the fully-updated system.
+    [switch] $CleanupOnly,
+
     # Administrator password — used only to write DefaultPassword for autologin.
     # Optional: if empty, DefaultPassword registry value is left as-is (Deploy may have set it).
     [string] $AdminPassword = '',
@@ -231,6 +236,97 @@ function Assert-Hash {
     Write-Host "  [OK] SHA256 verified: $Label" -ForegroundColor Green
 }
 
+function Invoke-DiskCleanup {
+    Write-Step "Cleaning up disk"
+
+    # 1. Windows Disk Cleanup (cleanmgr) — run with /sagerun:1 using preset flags
+    #    First register the flags via /sageset:1 in the registry (silent, no UI)
+    $cleanKey = 'HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\Explorer\VolumeCaches'
+    $sageset  = 1
+    $cleanCategories = @(
+        'Active Setup Temp Folders',
+        'BranchCache',
+        'Downloaded Program Files',
+        'Internet Cache Files',
+        'Memory Dump Files',
+        'Old ChkDsk Files',
+        'Previous Installations',
+        'Recycle Bin',
+        'Service Pack Cleanup',
+        'Setup Log Files',
+        'System error memory dump files',
+        'System error minidump files',
+        'Temporary Files',
+        'Temporary Setup Files',
+        'Thumbnail Cache',
+        'Update Cleanup',
+        'Upgrade Discarded Files',
+        'Windows Defender',
+        'Windows Error Reporting Archive Files',
+        'Windows Error Reporting Files',
+        'Windows Error Reporting Queue Files',
+        'Windows Error Reporting Temp Files',
+        'Windows ESD installation files',
+        'Windows Upgrade Log Files'
+    )
+    foreach ($cat in $cleanCategories) {
+        $keyPath = "$cleanKey\$cat"
+        if (Test-Path $keyPath) {
+            Set-ItemProperty -Path $keyPath -Name "StateFlags$('{0:D4}' -f $sageset)" -Value 2 -Type DWord -Force -ErrorAction SilentlyContinue
+        }
+    }
+    Write-Host "Running cleanmgr /sagerun:$sageset (may take a few minutes)..."
+    $proc = Start-Process -FilePath 'cleanmgr.exe' -ArgumentList "/sagerun:$sageset" -Wait -PassThru
+    Write-Host "cleanmgr exited (code $($proc.ExitCode))."
+
+    # 2. Clear Windows Update cache
+    Write-Host "Stopping WU-adjacent services before cache clear..."
+    foreach ($svc in 'wuauserv', 'UsoSvc', 'bits', 'dosvc', 'TrustedInstaller') {
+        Stop-Service -Name $svc -Force -ErrorAction SilentlyContinue
+    }
+    $sdDownload = 'C:\Windows\SoftwareDistribution\Download'
+    Remove-Item $sdDownload -Recurse -Force -ErrorAction SilentlyContinue
+    New-Item -ItemType Directory -Path $sdDownload -Force | Out-Null
+    Write-Host "Windows Update download cache cleared."
+
+    # 3. Clear CBS / component store logs
+    Remove-Item 'C:\Windows\Logs\CBS\*' -Force -ErrorAction SilentlyContinue
+
+    # 4. Clear TEMP folders
+    Remove-Item 'C:\Windows\Temp\*'          -Recurse -Force -ErrorAction SilentlyContinue
+    Remove-Item "$env:TEMP\*"                -Recurse -Force -ErrorAction SilentlyContinue
+
+    # 5. Clear Prefetch
+    Remove-Item 'C:\Windows\Prefetch\*'      -Force -ErrorAction SilentlyContinue
+
+    # 6. Run DISM component store cleanup (removes superseded update components)
+    Write-Host "Running DISM component store cleanup (StartComponentCleanup)..."
+    $dism = Start-Process -FilePath 'dism.exe' `
+        -ArgumentList '/Online','/Cleanup-Image','/StartComponentCleanup','/ResetBase' `
+        -Wait -PassThru
+    Write-Host "DISM exited (code $($dism.ExitCode))."
+
+    Write-Host "Disk cleanup complete."
+
+    # DISM / WaaSMedicSvc may reset WU service StartType — re-enforce Disabled.
+    foreach ($svc in 'wuauserv', 'UsoSvc') {
+        Set-Service -Name $svc -StartupType Disabled -ErrorAction SilentlyContinue
+    }
+    Set-ItemProperty -Path 'HKLM:\SYSTEM\CurrentControlSet\Services\wuauserv' `
+        -Name Start -Value 4 -Type DWord -Force -ErrorAction SilentlyContinue
+
+    # Report remaining WU cache size (WaaSMedicSvc restores files — informational only).
+    $sdFiles = Get-ChildItem 'C:\Windows\SoftwareDistribution\Download' -Recurse -File -Force -ErrorAction SilentlyContinue
+    $sdBytes = [long]0
+    if ($sdFiles) { foreach ($f in $sdFiles) { $sdBytes += $f.Length } }
+    $sdMB = [math]::Round($sdBytes / 1MB, 1)
+    Write-Host "  SoftwareDistribution\Download: $(@($sdFiles).Count) file(s), ${sdMB} MB remaining (WaaSMedicSvc best-effort, WU is disabled)."
+
+    Assert-Step 'Cleanup' 'wuauserv StartType Disabled after cache clear' {
+        (Get-Service wuauserv -ErrorAction Stop).StartType -eq 'Disabled'
+    }
+}
+
 # ── Pinned SHA256 hashes for downloaded installers (§1.3) ─────────────────────
 # Fill in the hash values below. Get them by:
 #   (Invoke-WebRequest '<URL>' -UseBasicParsing -OutFile 'tmp.bin'; (Get-FileHash 'tmp.bin' -Algorithm SHA256).Hash)
@@ -272,6 +368,20 @@ $script:Hashes = @{
     # vcpkg: cloned from GitHub via TLS; pin integrity with $VcpkgRef (commit SHA)
 }
 
+# ── CleanupOnly early-exit ────────────────────────────────────────────────────
+# Called by Prepare-WinBuild2025.ps1 after the Windows Update loop completes.
+# Skips all provisioning steps; runs only disk cleanup then exits.
+if ($CleanupOnly) {
+    if ($SkipCleanup) {
+        Write-Host "[Setup] -CleanupOnly + -SkipCleanup: nothing to do." -ForegroundColor Yellow
+        exit 0
+    }
+    Write-Host "[Setup] CleanupOnly mode — running disk cleanup on fully-updated system." -ForegroundColor Cyan
+    Invoke-DiskCleanup
+    Write-Host "[Setup] CleanupOnly complete." -ForegroundColor Green
+    exit 0
+}
+
 # ── Pre-flight: remove temp files from any previous partial run ───────────────
 # These files are normally deleted at the end of each step that creates them.
 # A previous interrupted run may have left them behind — remove before starting.
@@ -1527,114 +1637,9 @@ Write-Host "Tier-2 Toolchain installation complete." -ForegroundColor Green
 # ── Cleanup ───────────────────────────────────────────────────────────────────
 if ($SkipCleanup) {
     Write-Host "`n[Setup] Skipping disk cleanup (-SkipCleanup switch set)." -ForegroundColor Yellow
+} else {
+    Invoke-DiskCleanup
 }
-else {
-Write-Step "Cleaning up disk"
-
-# 1. Windows Disk Cleanup (cleanmgr) — run with /sagerun:1 using preset flags
-#    First register the flags via /sageset:1 in the registry (silent, no UI)
-$cleanKey = 'HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\Explorer\VolumeCaches'
-$sageset  = 1
-$cleanCategories = @(
-    'Active Setup Temp Folders',
-    'BranchCache',
-    'Downloaded Program Files',
-    'Internet Cache Files',
-    'Memory Dump Files',
-    'Old ChkDsk Files',
-    'Previous Installations',
-    'Recycle Bin',
-    'Service Pack Cleanup',
-    'Setup Log Files',
-    'System error memory dump files',
-    'System error minidump files',
-    'Temporary Files',
-    'Temporary Setup Files',
-    'Thumbnail Cache',
-    'Update Cleanup',
-    'Upgrade Discarded Files',
-    'Windows Defender',
-    'Windows Error Reporting Archive Files',
-    'Windows Error Reporting Files',
-    'Windows Error Reporting Queue Files',
-    'Windows Error Reporting Temp Files',
-    'Windows ESD installation files',
-    'Windows Upgrade Log Files'
-)
-foreach ($cat in $cleanCategories) {
-    $keyPath = "$cleanKey\$cat"
-    if (Test-Path $keyPath) {
-        Set-ItemProperty -Path $keyPath -Name "StateFlags$('{0:D4}' -f $sageset)" -Value 2 -Type DWord -Force -ErrorAction SilentlyContinue
-    }
-}
-Write-Host "Running cleanmgr /sagerun:$sageset (may take a few minutes)..."
-$proc = Start-Process -FilePath 'cleanmgr.exe' -ArgumentList "/sagerun:$sageset" -Wait -PassThru
-Write-Host "cleanmgr exited (code $($proc.ExitCode))."
-
-# 2. Clear Windows Update cache
-# Stop all services that might hold file locks inside SoftwareDistribution:
-#   wuauserv   — Windows Update (already Disabled from Step 6b, stop is idempotent)
-#   UsoSvc     — Update Orchestrator (already Disabled from Step 6b)
-#   bits       — Background Intelligent Transfer Service (downloads WU payloads)
-#   dosvc      — Delivery Optimization (peer-to-peer WU cache)
-#   TrustedInstaller — Windows Modules Installer (unpacks .cab/.msu update packages)
-Write-Host "Stopping WU-adjacent services before cache clear..."
-foreach ($svc in 'wuauserv', 'UsoSvc', 'bits', 'dosvc', 'TrustedInstaller') {
-    Stop-Service -Name $svc -Force -ErrorAction SilentlyContinue
-}
-
-# Delete and recreate the Download folder — more reliable than 'Remove-Item *'
-# because locked files inside subdirectories silently fail with wildcard delete.
-$sdDownload = 'C:\Windows\SoftwareDistribution\Download'
-Remove-Item $sdDownload -Recurse -Force -ErrorAction SilentlyContinue
-New-Item -ItemType Directory -Path $sdDownload -Force | Out-Null
-Write-Host "Windows Update download cache cleared."
-
-# 3. Clear CBS / component store logs
-Remove-Item 'C:\Windows\Logs\CBS\*' -Force -ErrorAction SilentlyContinue
-
-# 4. Clear TEMP folders
-Remove-Item 'C:\Windows\Temp\*'          -Recurse -Force -ErrorAction SilentlyContinue
-Remove-Item "$env:TEMP\*"                -Recurse -Force -ErrorAction SilentlyContinue
-
-# 5. Clear Prefetch
-Remove-Item 'C:\Windows\Prefetch\*'      -Force -ErrorAction SilentlyContinue
-
-# 6. Run DISM component store cleanup (removes superseded update components)
-Write-Host "Running DISM component store cleanup (StartComponentCleanup)..."
-$dism = Start-Process -FilePath 'dism.exe' `
-    -ArgumentList '/Online','/Cleanup-Image','/StartComponentCleanup','/ResetBase' `
-    -Wait -PassThru
-Write-Host "DISM exited (code $($dism.ExitCode))."
-
-Write-Host "Disk cleanup complete."
-
-# DISM StartComponentCleanup / WaaSMedicSvc may reset service StartType during cleanup.
-# Re-enforce Disabled on both WU services after DISM completes.
-foreach ($svc in 'wuauserv', 'UsoSvc') {
-    Set-Service -Name $svc -StartupType Disabled -ErrorAction SilentlyContinue
-}
-Set-ItemProperty -Path 'HKLM:\SYSTEM\CurrentControlSet\Services\wuauserv' `
-    -Name Start -Value 4 -Type DWord -Force -ErrorAction SilentlyContinue
-
-# Validation — leftover artifacts from this run should be gone
-# Note: do NOT assert SoftwareDistribution\Download empty.
-# WaaSMedicSvc (Windows Update Medic Service) is tamper-protected — cannot be stopped —
-# and actively restores WU files immediately after deletion. Files reappearing here is
-# normal and harmless: WU is permanently disabled by Step 6b (service Disabled + GPO keys).
-# Report remaining size as informational only.
-$sdFiles = Get-ChildItem 'C:\Windows\SoftwareDistribution\Download' -Recurse -File -Force -ErrorAction SilentlyContinue
-# Avoid Measure-Object .Sum: under Set-StrictMode -Version Latest, PS 5.1 treats the
-# nullable Sum property as absent when the pipe is empty, throwing "property not found".
-$sdBytes = [long]0
-if ($sdFiles) { foreach ($f in $sdFiles) { $sdBytes += $f.Length } }
-$sdMB = [math]::Round($sdBytes / 1MB, 1)
-Write-Host "  SoftwareDistribution\Download: $(@($sdFiles).Count) file(s), ${sdMB} MB remaining (WaaSMedicSvc best-effort, WU is disabled)."
-
-Assert-Step 'Cleanup' 'wuauserv StartType Disabled after cache clear' {
-    (Get-Service wuauserv -ErrorAction Stop).StartType -eq 'Disabled'
-}
-} # end if (-not $SkipCleanup)
 
 # ── Pre-Final re-affirmation ──────────────────────────────────────────────────
 # Steps 7-10 (installs, cleanup, DISM) can take hours. During that window background

tests/python/test_commands_retention.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 import shutil
 import time
 from pathlib import Path
-from typing import Any
 
 import pytest
 from click.testing import CliRunner
@@ -13,7 +12,6 @@ from click.testing import CliRunner
 import ci_orchestrator.commands.retention as ret
 from ci_orchestrator.__main__ import cli
 
-
 # ── helpers ───────────────────────────────────────────────────────────────────
 
 

tests/python/test_commands_template.py

@@ -17,7 +17,6 @@ import ci_orchestrator.commands.template as tmpl_mod
 from ci_orchestrator.__main__ import cli
 from ci_orchestrator.transport.winrm import WinRmResult
 
-
 # ── fixtures ──────────────────────────────────────────────────────────────────
 
 
@@ -345,7 +344,7 @@ def test_parse_exit_marker_malformed() -> None:
 
 def test_wait_tcp_success(monkeypatch: pytest.MonkeyPatch) -> None:
     class _FakeConn:
-        def __enter__(self) -> "_FakeConn":
+        def __enter__(self) -> _FakeConn:
             return self
 
         def __exit__(self, *_: object) -> None:
@@ -699,7 +698,7 @@ def test_prepare_win_recreate_snapshot_flag(
 def test_prepare_win_skip_flags(
     tmp_path: Path, monkeypatch: pytest.MonkeyPatch
 ) -> None:
-    """--skip-windows-update/tier2/cleanup flags reach the invoke script."""
+    """-SkipCleanup:$true always in loop args; --skip-cleanup suppresses post-loop cleanup call."""
     vmx = _setup_prepare_env(tmp_path, with_static_ip=False)
     monkeypatch.chdir(tmp_path)
     backend = _make_backend([True, False])
@@ -1030,3 +1029,116 @@ def test_prepare_win_store_credential_default_target(
     assert result.exit_code == 0, result.output
     assert any(s[0] == "BuildVMGuest:meta" for s in stored)
     assert any(s[0] == "BuildVMGuest" for s in stored)
+
+
+def test_prepare_win_cleanup_runs_after_loop(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """CleanupOnly invocation fires after WU loop exits 0 (default, no --skip-cleanup)."""
+    vmx = _setup_prepare_env(tmp_path, with_static_ip=False)
+    monkeypatch.chdir(tmp_path)
+    backend = _make_backend([True, False])
+
+    captured_scripts: list[str] = []
+
+    def _run(script: str = "", **kwargs: object) -> WinRmResult:
+        captured_scripts.append(script)
+        if "ConvertTo-Json" in script:
+            return WinRmResult(stdout="{}", stderr="", returncode=0)
+        return WinRmResult(stdout="CI_EXITCODE:0", stderr="", returncode=0)
+
+    transport: MagicMock = MagicMock()
+    transport.run.side_effect = _run
+    transport.is_ready.return_value = True
+
+    _patch_prepare(monkeypatch, backend, transport)
+
+    result = CliRunner().invoke(
+        cli,
+        [
+            "template", "prepare-win",
+            "--vmx", str(vmx),
+            "--ip", "192.168.1.100",
+            "--admin-password", "adminpass",
+            "--build-password", "buildpass",
+        ],
+    )
+    assert result.exit_code == 0, result.output
+    assert any("CleanupOnly:$true" in s for s in captured_scripts), (
+        "Expected a -CleanupOnly:$true script after the WU loop"
+    )
+    assert "Post-update cleanup complete" in result.output
+
+
+def test_prepare_win_skip_cleanup_no_cleanup_call(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """--skip-cleanup suppresses the post-loop CleanupOnly invocation."""
+    vmx = _setup_prepare_env(tmp_path, with_static_ip=False)
+    monkeypatch.chdir(tmp_path)
+    backend = _make_backend([True, False])
+
+    captured_scripts: list[str] = []
+
+    def _run(script: str = "", **kwargs: object) -> WinRmResult:
+        captured_scripts.append(script)
+        if "ConvertTo-Json" in script:
+            return WinRmResult(stdout="{}", stderr="", returncode=0)
+        return WinRmResult(stdout="CI_EXITCODE:0", stderr="", returncode=0)
+
+    transport: MagicMock = MagicMock()
+    transport.run.side_effect = _run
+    transport.is_ready.return_value = True
+
+    _patch_prepare(monkeypatch, backend, transport)
+
+    result = CliRunner().invoke(
+        cli,
+        [
+            "template", "prepare-win",
+            "--vmx", str(vmx),
+            "--ip", "192.168.1.100",
+            "--admin-password", "adminpass",
+            "--build-password", "buildpass",
+            "--skip-cleanup",
+        ],
+    )
+    assert result.exit_code == 0, result.output
+    assert not any("CleanupOnly:$true" in s for s in captured_scripts), (
+        "--skip-cleanup must suppress the -CleanupOnly:$true call"
+    )
+
+
+def test_prepare_win_cleanup_only_fails(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Cleanup-only step returning nonzero exits the command with an error."""
+    vmx = _setup_prepare_env(tmp_path, with_static_ip=False)
+    monkeypatch.chdir(tmp_path)
+    backend = _make_backend([True])
+
+    def _run(script: str = "", **kwargs: object) -> WinRmResult:
+        if "CleanupOnly" in script:
+            return WinRmResult(stdout="CI_EXITCODE:1", stderr="cleanup error", returncode=0)
+        if "ConvertTo-Json" in script:
+            return WinRmResult(stdout="{}", stderr="", returncode=0)
+        return WinRmResult(stdout="CI_EXITCODE:0", stderr="", returncode=0)
+
+    transport: MagicMock = MagicMock()
+    transport.run.side_effect = _run
+    transport.is_ready.return_value = True
+
+    _patch_prepare(monkeypatch, backend, transport)
+
+    result = CliRunner().invoke(
+        cli,
+        [
+            "template", "prepare-win",
+            "--vmx", str(vmx),
+            "--ip", "192.168.1.100",
+            "--admin-password", "adminpass",
+            "--build-password", "buildpass",
+        ],
+    )
+    assert result.exit_code != 0
+    assert "Disk cleanup failed" in result.output