Flamenco-Management/remote_worker_controller.ps1

param(
    [Parameter(Mandatory = $true)]
    [string]$WorkerName,

    [Parameter(Mandatory = $true)]
    [string]$WorkerType,

    [Parameter(Mandatory = $true)]
    [string]$PayloadBase64,

[string]$DataRoot = (Join-Path ([Environment]::GetFolderPath('LocalApplicationData')) 'UnifiedWorkers'),

    [int]$MaxRestarts = 5,

    [int]$RestartDelaySeconds = 10
)

$ErrorActionPreference = 'Stop'

# region Path setup
$workerRoot = Join-Path -Path $DataRoot -ChildPath $WorkerType
$instanceRoot = Join-Path -Path $workerRoot -ChildPath $WorkerName
New-Item -ItemType Directory -Path $instanceRoot -Force | Out-Null

$logsRoot = Join-Path -Path $instanceRoot -ChildPath 'logs'
$stateRoot = Join-Path -Path $instanceRoot -ChildPath 'state'
New-Item -ItemType Directory -Path $logsRoot -Force | Out-Null
New-Item -ItemType Directory -Path $stateRoot -Force | Out-Null

$logPath = Join-Path -Path $logsRoot -ChildPath 'worker.log'
$metaPath = Join-Path -Path $stateRoot -ChildPath 'worker-info.json'
$commandPath = Join-Path -Path $stateRoot -ChildPath 'commands.txt'
$payloadPath = Join-Path -Path $stateRoot -ChildPath "payload-$([Guid]::NewGuid().ToString()).ps1"
# endregion

# region Helpers
function Write-ControllerLog {
    param([string]$Message)
    $timestamp = (Get-Date).ToString('u')
    Add-Content -Path $logPath -Value "[CTRL $timestamp] $Message" -Encoding UTF8
}

function Write-Metadata {
    param(
        [string]$Status,
        [nullable[int]]$WorkerPid = $null,
        [nullable[int]]$ControllerPid = $PID,
        [int]$Restarts = 0,
        [nullable[int]]$LastExitCode = $null
    )

    $payload = [pscustomobject]@{
        WorkerName     = $WorkerName
        WorkerType     = $WorkerType
        Status         = $Status
        ControllerPid  = $ControllerPid
        WorkerPid      = $WorkerPid
        Restarts       = $Restarts
        LastExitCode   = $LastExitCode
        LogPath        = $logPath
        CommandPath    = $commandPath
        PayloadPath    = $payloadPath
        UpdatedAtUtc   = (Get-Date).ToUniversalTime()
    }

    $payload | ConvertTo-Json -Depth 5 | Set-Content -Path $metaPath -Encoding UTF8
}

function Get-PendingCommands {
    if (-not (Test-Path $commandPath)) {
        return @()
    }

    try {
        $lines = Get-Content -Path $commandPath -ErrorAction Stop
        Remove-Item -Path $commandPath -Force -ErrorAction SilentlyContinue
        return $lines | Where-Object { -not [string]::IsNullOrWhiteSpace($_) }
    }
    catch {
        return @()
    }
}
# endregion

try {
    # Write payload script to disk
    $payloadBytes = [Convert]::FromBase64String($PayloadBase64)
    [IO.File]::WriteAllBytes($payloadPath, $payloadBytes)
    Write-ControllerLog "Payload written to $payloadPath"
}
catch {
    Write-Error "Unable to write payload: $($_.Exception.Message)"
    exit 1
}

$restartCount = 0
$controllerPid = $PID

while ($restartCount -le $MaxRestarts) {
    try {
        # Initialize worker process
        $psi = [System.Diagnostics.ProcessStartInfo]::new()
        $pwsh = Get-Command pwsh -ErrorAction SilentlyContinue
        if ($pwsh) {
            $psi.FileName = $pwsh.Source
        }
        else {
            $psi.FileName = (Get-Command powershell -ErrorAction Stop).Source
        }

        $psi.Arguments = "-NoLogo -NoProfile -ExecutionPolicy Bypass -File `"$payloadPath`""
        $psi.UseShellExecute = $false
        $psi.RedirectStandardInput = $true
        $psi.RedirectStandardOutput = $true
        $psi.RedirectStandardError = $true
        $psi.CreateNoWindow = $true

        $workerProcess = New-Object System.Diagnostics.Process
        $workerProcess.StartInfo = $psi

        if (-not $workerProcess.Start()) {
            throw "Failed to start worker process."
        }

        Write-ControllerLog "Worker process started with PID $($workerProcess.Id)"
        Write-Metadata -Status 'running' -WorkerPid $workerProcess.Id -ControllerPid $controllerPid -Restarts $restartCount

        $logWriter = [System.IO.StreamWriter]::new($logPath, $true, [System.Text.Encoding]::UTF8)
        $logWriter.AutoFlush = $true

        $outputHandler = [System.Diagnostics.DataReceivedEventHandler]{
            param($s, $e)
            if ($e.Data) {
                $ts = (Get-Date).ToString('u')
                $logWriter.WriteLine("[OUT $ts] $($e.Data)")
            }
        }
        $errorHandler = [System.Diagnostics.DataReceivedEventHandler]{
            param($s, $e)
            if ($e.Data) {
                $ts = (Get-Date).ToString('u')
                $logWriter.WriteLine("[ERR $ts] $($e.Data)")
            }
        }

        $workerProcess.add_OutputDataReceived($outputHandler)
        $workerProcess.add_ErrorDataReceived($errorHandler)
        $workerProcess.BeginOutputReadLine()
        $workerProcess.BeginErrorReadLine()

        while (-not $workerProcess.HasExited) {
            $commands = Get-PendingCommands
            foreach ($command in $commands) {
                $trimmed = $command.Trim()
                if (-not $trimmed) { continue }

                Write-ControllerLog "Received command '$trimmed'"
                try {
                    $workerProcess.StandardInput.WriteLine($trimmed)
                    $workerProcess.StandardInput.Flush()
                }
                catch {
                    Write-ControllerLog "Failed to forward command '$trimmed': $($_.Exception.Message)"
                }

                if ($trimmed -ieq 'quit') {
                    Write-ControllerLog "Quit command issued. Waiting for worker to exit."
                }
            }

            Start-Sleep -Milliseconds 500
        }

        $exitCode = $workerProcess.ExitCode
        Write-ControllerLog "Worker exited with code $exitCode"
        Write-Metadata -Status 'stopped' -WorkerPid $null -ControllerPid $controllerPid -Restarts $restartCount -LastExitCode $exitCode

        if ($exitCode -eq 0) {
            break
        }
    }
    catch {
        Write-ControllerLog "Controller error: $($_.Exception.Message)"
    }

    $restartCount++
    if ($restartCount -gt $MaxRestarts) {
        Write-ControllerLog "Maximum restart attempts reached. Controller stopping."
        break
    }

    Write-ControllerLog "Restarting worker in $RestartDelaySeconds seconds (attempt $restartCount of $MaxRestarts)."
    Start-Sleep -Seconds $RestartDelaySeconds
}

Write-Metadata -Status 'inactive' -WorkerPid $null -ControllerPid $controllerPid -Restarts $restartCount
Write-ControllerLog "Controller exiting."