fix: validate Windows node reset task result

parts/windows/kuberneteswindowssetup.ps1

@@ -651,21 +651,7 @@ function NodePrep {
         Postpone-RestartComputer
     } else {
         Logs-To-Event -TaskName "AKS.WindowsCSE.StartScheduledTask" -TaskMessage "Setup Complete, start NodeResetScriptTask to register Windows node without reboot"
-        Start-ScheduledTask -TaskName "k8s-restart-job"
-
-        $timeout = 180 ##  seconds
-        $timer = [Diagnostics.Stopwatch]::StartNew()
-        while ((Get-ScheduledTask -TaskName 'k8s-restart-job').State -ne 'Ready') {
-            # The task `k8s-restart-job` needs ~8 seconds.
-            if ($timer.Elapsed.TotalSeconds -gt $timeout) {
-                Set-ExitCode -ExitCode $global:WINDOWS_CSE_ERROR_START_NODE_RESET_SCRIPT_TASK -ErrorMessage "NodeResetScriptTask is not finished after [$($timer.Elapsed.TotalSeconds)] seconds"
-            }
-
-            Write-Log -Message "Waiting on NodeResetScriptTask..."
-            Start-Sleep -Seconds 3
-        }
-        $timer.Stop()
-        Write-Log -Message "We waited [$($timer.Elapsed.TotalSeconds)] seconds on NodeResetScriptTask"
+        Start-NodeResetScriptTask
     }
     Write-Log "NodePrep completed successfully"
     Logs-To-Event -TaskName "AKS.WindowsCSE.NodePrep" -TaskMessage "NodePrep completed successfully"

parts/windows/windowscsehelper.ps1

@@ -300,6 +300,50 @@ function Set-ExitCode
     exit $ExitCode
 }
 
+function Start-NodeResetScriptTask
+{
+    Param(
+        [Parameter(Mandatory=$false)][int]
+        $TimeoutSeconds = 180
+    )
+
+    $taskName = "k8s-restart-job"
+    $taskRunningResult = 0x00041301
+    $previousRunTime = (Get-ScheduledTaskInfo -TaskName $taskName).LastRunTime
+    Start-ScheduledTask -TaskName $taskName
+
+    $timer = [Diagnostics.Stopwatch]::StartNew()
+    do {
+        $taskInfo = Get-ScheduledTaskInfo -TaskName $taskName
+        $task = Get-ScheduledTask -TaskName $taskName
+        if ($task.State -eq "Ready" -and $taskInfo.LastRunTime -ne $previousRunTime) {
+            $taskInfo = Get-ScheduledTaskInfo -TaskName $taskName
+            if ($taskInfo.LastRunTime -ne $previousRunTime -and $taskInfo.LastTaskResult -ne $taskRunningResult) {
+                break
+            }
+        }
+
+        if ($timer.Elapsed.TotalSeconds -gt $TimeoutSeconds) {
+            Set-ExitCode -ExitCode $global:WINDOWS_CSE_ERROR_START_NODE_RESET_SCRIPT_TASK -ErrorMessage "NodeResetScriptTask is not finished after [$($timer.Elapsed.TotalSeconds)] seconds"
+        }
+
+        Write-Log -Message "Waiting on NodeResetScriptTask..."
+        Start-Sleep -Seconds 3
+    } while ($true)
+    $timer.Stop()
+
+    if ($taskInfo.LastTaskResult -ne 0) {
+        Set-ExitCode -ExitCode $global:WINDOWS_CSE_ERROR_START_NODE_RESET_SCRIPT_TASK -ErrorMessage "NodeResetScriptTask failed with result $($taskInfo.LastTaskResult)"
+    }
+
+    $kubeletService = Get-Service -Name "kubelet" -ErrorAction SilentlyContinue
+    if ($null -eq $kubeletService -or $kubeletService.Status -ne "Running") {
+        Set-ExitCode -ExitCode $global:WINDOWS_CSE_ERROR_START_NODE_RESET_SCRIPT_TASK -ErrorMessage "kubelet service is not running after NodeResetScriptTask completed"
+    }
+
+    Write-Log -Message "We waited [$($timer.Elapsed.TotalSeconds)] seconds on NodeResetScriptTask"
+}
+
 function Postpone-RestartComputer
 {
     Logs-To-Event -TaskName "AKS.WindowsCSE.PostponeRestartComputer" -TaskMessage "Start to create an one-time task to restart the VM"

parts/windows/windowscsehelper.tests.ps1

@@ -589,3 +589,81 @@ Describe "Update-BaseUrl" {
     $result | Should -Be "https://packages.aks.azure.com/path/to/resource"
   }
 }
+
+Describe "Start-NodeResetScriptTask" {
+  BeforeEach {
+    $script:taskInfoCallCount = 0
+    Mock Start-ScheduledTask -MockWith {}
+    Mock Get-ScheduledTask -MockWith { return [pscustomobject]@{ State = "Ready" } }
+    Mock Get-ScheduledTaskInfo -MockWith {
+      $script:taskInfoCallCount++
+      if ($script:taskInfoCallCount -eq 1) {
+        return [pscustomobject]@{ LastRunTime = [datetime]"2026-01-01"; LastTaskResult = 0 }
+      }
+      return [pscustomobject]@{ LastRunTime = [datetime]"2026-01-02"; LastTaskResult = 0 }
+    }
+    Mock Get-Service -MockWith { return [pscustomobject]@{ Status = "Running" } }
+    Mock Start-Sleep -MockWith {}
+    Mock Write-Log -MockWith {}
+    Mock Set-ExitCode -MockWith {
+      Param($ExitCode, $ErrorMessage)
+      throw $ErrorMessage
+    }
+  }
+
+  It "waits for a new successful run and running kubelet" {
+    Start-NodeResetScriptTask
+
+    Assert-MockCalled -CommandName Start-ScheduledTask -Exactly -Times 1
+    Assert-MockCalled -CommandName Set-ExitCode -Exactly -Times 0
+  }
+
+  It "does not accept Ready before the new run starts" {
+    Mock Get-ScheduledTaskInfo -MockWith {
+      $script:taskInfoCallCount++
+      if ($script:taskInfoCallCount -le 2) {
+        return [pscustomobject]@{ LastRunTime = [datetime]"2026-01-01"; LastTaskResult = 0 }
+      }
+      return [pscustomobject]@{ LastRunTime = [datetime]"2026-01-02"; LastTaskResult = 0 }
+    }
+
+    Start-NodeResetScriptTask
+
+    Assert-MockCalled -CommandName Start-Sleep -Exactly -Times 1
+  }
+
+  It "does not accept stale Ready while the task is running" {
+    Mock Get-ScheduledTaskInfo -MockWith {
+      $script:taskInfoCallCount++
+      if ($script:taskInfoCallCount -eq 1) {
+        return [pscustomobject]@{ LastRunTime = [datetime]"2026-01-01"; LastTaskResult = 0 }
+      }
+      if ($script:taskInfoCallCount -le 3) {
+        return [pscustomobject]@{ LastRunTime = [datetime]"2026-01-02"; LastTaskResult = 0x00041301 }
+      }
+      return [pscustomobject]@{ LastRunTime = [datetime]"2026-01-02"; LastTaskResult = 0 }
+    }
+
+    Start-NodeResetScriptTask
+
+    Assert-MockCalled -CommandName Start-Sleep -Exactly -Times 1
+  }
+
+  It "fails when the task result is nonzero" {
+    Mock Get-ScheduledTaskInfo -MockWith {
+      $script:taskInfoCallCount++
+      if ($script:taskInfoCallCount -eq 1) {
+        return [pscustomobject]@{ LastRunTime = [datetime]"2026-01-01"; LastTaskResult = 0 }
+      }
+      return [pscustomobject]@{ LastRunTime = [datetime]"2026-01-02"; LastTaskResult = 1 }
+    }
+
+    { Start-NodeResetScriptTask } | Should -Throw "*failed with result 1*"
+  }
+
+  It "fails when kubelet is not running" {
+    Mock Get-Service -MockWith { return [pscustomobject]@{ Status = "Stopped" } }
+
+    { Start-NodeResetScriptTask } | Should -Throw "*kubelet service is not running*"
+  }
+}