ScriptBeacon.psm1

#requires -Version 5.1
Set-StrictMode -Version Latest
try { [System.Net.ServicePointManager]::SecurityProtocol = [System.Net.SecurityProtocolType]::Tls12 } catch {}

# ------------------------ Module State ------------------------
$script:BeamState = [ordered]@{
    IsOpen              = $false
    NestDepth           = 0
    AggregatePerProcess = $true

    TranscriptPath = $null   # Raw transcript (as produced by Start-Transcript)
    CleanPath      = $null   # Cleaned log (no transcript header/footer)
    GzipPath       = $null

    StartedAtUtc = $null
    EndedAtUtc   = $null
    Tags         = @{}

    BleepId     = $env:BLEEP_ID
    ApiRoot     = if ($env:BLEEP_API) { $env:BLEEP_API } else { 'https://api.scriptbeacon.com' }
    WriteSecret = $env:BEACON_WRITE

    ContentType     = 'text/plain; charset=utf-8'
    ContentEncoding = 'gzip'

    RunId      = $null
    UploadId   = $null
    ObjectKey  = $null
    Finalized  = $false
    ExitCode   = 0

    # Crash guards registration flags
    ProcessExitHookRegistered  = $false
    UnhandledHookRegistered    = $false

    # Hold a preformatted host-like error block to append post-transcript
    PendingErrorBlock = $null

    # Optional forced paths (override autodiscovery)
    InitPathOverride     = $null # e.g. "/b/{id}/chunk-init"
    CompletePathOverride = $null # e.g. "/b/{id}/chunk-complete"
    IngestPathOverride   = $null # e.g. "/b/{id}" or "/lb/{id}"

    # control how ScriptStackTrace is included in the appended error block
    # values: 'none' (default), 'user', 'full'
    StackTraceMode = 'none'
}

# Candidate endpoint matrices (most likely first)
$script:InitCandidates = @(
    "/b/{id}/chunk-init"
)
$script:CompleteCandidates = @(
    "/b/{id}/chunk-complete"
)
$script:IngestCandidates = @(
    "/b/{id}"
)

# -------- Public configuration --------
function Set-BeaconConfig {
    [CmdletBinding()]
    param(
        [Parameter(Position=0)] [ValidateNotNullOrEmpty()] [string] $Id,
        [Parameter()]            [ValidateNotNullOrEmpty()] [string] $Api = $null,
        [Parameter()]            [string] $WriteSecret = $null,
        [switch] $NoAggregate,
        # Optional: override backend routes
        [string] $InitPath,
        [string] $CompletePath,
        [string] $IngestPath,
        # control appended stack trace detail
        [ValidateSet('none','user','full')] [string] $StackTrace
    )
    if ($PSBoundParameters.ContainsKey('Id'))          { $script:BeamState.BleepId = $Id }
    if ($PSBoundParameters.ContainsKey('Api'))         { $script:BeamState.ApiRoot = $Api.TrimEnd('/') }
    if ($PSBoundParameters.ContainsKey('WriteSecret')) { $script:BeamState.WriteSecret = $WriteSecret }
    if ($NoAggregate.IsPresent) { $script:BeamState.AggregatePerProcess = $false }

    if ($PSBoundParameters.ContainsKey('InitPath'))     { $script:BeamState.InitPathOverride     = $InitPath }
    if ($PSBoundParameters.ContainsKey('CompletePath')) { $script:BeamState.CompletePathOverride = $CompletePath }
    if ($PSBoundParameters.ContainsKey('IngestPath'))   { $script:BeamState.IngestPathOverride   = $IngestPath }

    if ($PSBoundParameters.ContainsKey('StackTrace'))   { $script:BeamState.StackTraceMode       = $StackTrace }
}

function Get-BeamStatus { [CmdletBinding()] param() [pscustomobject]$script:BeamState }

# -------- Helpers --------
function New-IsoNowUtc { (Get-Date).ToUniversalTime().ToString('yyyy-MM-ddTHH:mm:ssZ') }

function Merge-BeamTags {
    param([hashtable]$Base, [object]$Extra)
    $merged = @{}
    if ($Base) { $Base.GetEnumerator() | ForEach-Object { $merged[$_.Key] = "$($_.Value)" } }

    if ($Extra) {
        # NEW: single string "a=b, c=d; e=f" support
        if ($Extra -is [string]) {
            $pairs = $Extra -split '[,;]' | ForEach-Object { $_.Trim() } | Where-Object { $_ }
            foreach ($kv in $pairs) {
                if ($kv -is [string] -and $kv.Contains('=')) {
                    $parts = $kv.Split('=',2)
                    $k = $parts[0].Trim(); $v = if ($parts.Count -gt 1) { $parts[1].Trim() } else { '' }
                    if ($k) { $merged[$k] = $v }
                }
            }
        } elseif ($Extra -is [hashtable]) {
            $Extra.GetEnumerator() | ForEach-Object { $merged[$_.Key] = "$($_.Value)" }
        } elseif ($Extra -is [System.Collections.IEnumerable]) {
            foreach ($kv in $Extra) {
                if ($kv -is [string] -and $kv.Contains('=')) {
                    $parts = $kv.Split('=',2)
                    $k = $parts[0]; $v = if ($parts.Count -gt 1) { $parts[1] } else { '' }
                    if ($k) { $merged[$k] = $v }
                } elseif ($kv -is [System.Collections.DictionaryEntry]) {
                    $merged[$kv.Key] = "$($kv.Value)"
                }
            }
        }
    }

    # Standard enrichers
    $hostName = $env:COMPUTERNAME; if ([string]::IsNullOrEmpty($hostName)) { $hostName = [System.Environment]::MachineName }
    $shellName = if ($PSVersionTable.PSEdition -eq 'Core') { 'pwsh' } else { 'powershell' }
    $merged['host']  = $hostName
    $merged['user']  = [Environment]::UserName
    $merged['shell'] = "$shellName-" + $PSVersionTable.PSVersion.ToString()
    $merged
}

function Normalize-InitResponse {
    param([Parameter(Mandatory)][object]$Resp)
    if ($Resp.PSObject.Properties['data']) { $Resp = $Resp.data }
    if ($Resp -is [array]) { $Resp = $Resp[0] }
    $out = [ordered]@{ run_id=$null; upload_id=$null; upload_url=$null; object_key=$null; required_headers=@{} }
    $p = $Resp.PSObject.Properties
    foreach ($k in 'run_id','runId','id','session_run_id','sessionRunId') { if ($p[$k]) { $out.run_id = [string]$p[$k].Value; break } }
    foreach ($k in 'upload_id','uploadId','session_id','sessionId','upload_session_id','uploadSessionId') { if ($p[$k]) { $out.upload_id = [string]$p[$k].Value; break } }
    if (-not $out.run_id -and $p['run']   -and $Resp.run.PSObject.Properties['id']) { $out.run_id = [string]$Resp.run.id }
    if (-not $out.upload_id -and $p['upload'] -and $Resp.upload.PSObject.Properties['id']) { $out.upload_id = [string]$Resp.upload.id }
    foreach ($k in 'upload_url','uploadUrl','url','put_url','putUrl') { if ($p[$k]) { $out.upload_url = [string]$p[$k].Value; break } }
    foreach ($k in 'object_key','objectKey','key','blob_key','blobKey','s3_key') { if ($p[$k]) { $out.object_key = [string]$p[$k].Value; break } }
    $rh = $null; foreach ($k in 'required_headers','requiredHeaders','headers','signed_headers') { if ($p[$k]) { $rh = $p[$k].Value; break } }
    if ($rh) { $rh.PSObject.Properties | ForEach-Object { $out.required_headers[$_.Name] = [string]$_.Value } }
    $out
}

function Normalize-CompleteResponse {
    [CmdletBinding()] param([object]$Resp,[string]$FallbackRunId,[string]$FallbackUploadId,[Nullable[long]]$Bytes,[string]$ETag)
    if ($Resp -and $Resp.PSObject.Properties['data']) { $Resp = $Resp.data }
    if ($Resp -is [array]) { $Resp = $Resp[0] }
    $runId=$null;$upload=$null;$bytes=$null;$etag=$null
    if ($Resp) {
        $p = $Resp.PSObject.Properties
        foreach ($k in 'run_id','runId','id','session_run_id','sessionRunId') { if ($p[$k]) { $runId = [string]$p[$k].Value; break } }
        foreach ($k in 'upload_id','uploadId','session_id','sessionId','upload_session_id','uploadSessionId') { if ($p[$k]) { $upload = [string]$p[$k].Value; break } }
        foreach ($k in 'bytes','size_bytes','sizeBytes') { if ($p[$k]) { $bytes = [long]$p[$k].Value; break } }
        foreach ($k in 'etag','eTag','ETag') { if ($p[$k]) { $etag = [string]$p[$k].Value; break } }
    }
    if (-not $runId)  { $runId  = $FallbackRunId }
    if (-not $upload) { $upload = $FallbackUploadId }
    if (-not $bytes -and $Bytes) { $bytes = $Bytes }
    if (-not $etag  -and $ETag)  { $etag  = $ETag }
    [pscustomobject]@{ run_id=$runId; upload_id=$upload; bytes=$bytes; etag=$etag }
}

function Compress-Gzip {
    [CmdletBinding()] param(
        [Parameter(Mandatory)][ValidateScript({Test-Path $_})][string]$Path,
        [Parameter(Mandatory)][string]$Destination
    )
    $in=[System.IO.File]::OpenRead($Path)
    try {
        $out=[System.IO.File]::Create($Destination)
        try {
            $gzip=New-Object System.IO.Compression.GZipStream($out,[System.IO.Compression.CompressionLevel]::Optimal,$true)
            try { $in.CopyTo($gzip) } finally { $gzip.Dispose() }
        } finally { $out.Dispose() }
    } finally { $in.Dispose() }
}

# ---- Event subscription utility (prevents duplicate SUBSCRIBER_EXISTS noise)
function Ensure-EventSubscription {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)][string]$SourceIdentifier,
        [Parameter(Mandatory)][ScriptBlock]$RegisterAction
    )
    try { $existing = Get-EventSubscriber -SourceIdentifier $SourceIdentifier -ErrorAction SilentlyContinue } catch { $existing = $null }
    if ($existing) { return $true }
    try { & $RegisterAction; return $true } catch { if ($_.FullyQualifiedErrorId -like 'SUBSCRIBER_EXISTS*') { return $true } throw }
}

# ---- Build a host-like terminating-error block to append to transcript
function Format-ErrorRecordBlock {
    [CmdletBinding()]
    param([Parameter(Mandatory)][System.Management.Automation.ErrorRecord]$ErrorRecord)

    try {
        $inv     = $ErrorRecord.InvocationInfo
        $invName = $null
        if ($inv) { try { $invName = $inv.InvocationName } catch {} }
        $message = $ErrorRecord.Exception.Message

        $excType = $null
        try { if ($ErrorRecord.Exception) { $excType = $ErrorRecord.Exception.GetType().FullName } } catch {}

        $header = if ($invName) { "$invName : $message" } else { $message }

        $lines = New-Object System.Collections.Generic.List[string]
        $null  = $lines.Add($header)

        if ($excType) { $null = $lines.Add("Exception : $excType") }

        # Host-like position block ("At <path>:line char")
        $posMsg = $null
        try { if ($inv -and $inv.PositionMessage) { $posMsg = $inv.PositionMessage.TrimEnd() } } catch {}
        if ($posMsg) { $null = $lines.Add($posMsg) }

        $null = $lines.Add(" + CategoryInfo : $($ErrorRecord.CategoryInfo)")
        $null = $lines.Add(" + FullyQualifiedErrorId : $($ErrorRecord.FullyQualifiedErrorId)")

        # Optional ScriptStackTrace (configurable)
        $mode = $script:BeamState.StackTraceMode
        if (-not $mode) { $mode = 'none' }
        if ($mode -ne 'none') {
            $sst = $ErrorRecord.ScriptStackTrace
            if ($sst) {
                $stackLines = $sst -split "\r?\n"
                if ($mode -eq 'user') {
                    # keep only the first user .ps1 frame; drop module frames
                    $stackLines = $stackLines | Where-Object { $_ -match '\.ps1(:|,)' -and $_ -notmatch 'ScriptBeacon\.psm1' }
                    if ($stackLines.Count -gt 1) { $stackLines = @($stackLines[0]) }
                }
                if ($stackLines -and $stackLines.Count -gt 0) {
                    # print one label line + subsequent frames indented (if any)
                    $null = $lines.Add(" + ScriptStackTrace : $($stackLines[0])")
                    for ($si=1; $si -lt $stackLines.Count; $si++) {
                        $null = $lines.Add(" $($stackLines[$si])")
                    }
                }
            }
        }

        return ($lines -join [Environment]::NewLine)
    } catch {
        # Fallback: minimal line
        return $ErrorRecord.ToString()
    }
}

# ---- Network retry wrapper
function Invoke-WithRetry {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)][ScriptBlock]$Action,
        [int]$MaxAttempts = 3,
        [int]$BaseDelayMs = 300
    )
    $attempt = 0
    while ($true) {
        $attempt++
        try {
            return & $Action
        }
        catch {
            # Extract minimal HTTP context (status + Retry-After)
            $status = $null; $retryAfterS = $null
            try {
                $resp = $_.Exception.Response
                if ($resp) {
                    try { $status = [int]$resp.StatusCode } catch {}
                    try {
                        if ($resp.Headers) {
                            $retryAfterS = [string]$resp.Headers['Retry-After']
                            if (-not $retryAfterS -and $resp -is [System.Net.HttpWebResponse]) {
                                $retryAfterS = [string]$resp.GetResponseHeader('Retry-After')
                            }
                        }
                    } catch {}
                }
            } catch {}

            # Decide whether to retry
            $retryableStatuses = @(408, 429, 500, 502, 503, 504)
            $hasNumericRetryAfter = $false; $retryAfterMs = $null
            if ($retryAfterS) {
                $secVal = 0
                if ([int]::TryParse([string]$retryAfterS, [ref]$secVal) -and $secVal -gt 0) {
                    $hasNumericRetryAfter = $true
                    $retryAfterMs = [Math]::Min($secVal * 1000, 30000)
                }
            }

            # IMPORTANT tweak:
            # - 429 *without* Retry-After → treat as non‑retryable (per-beacon throttle returns 429 JSON but no header). :contentReference[oaicite:0]{index=0}
            # - Auth/guard rails (401/403) are not retryable.
            $shouldRetry =
                ($status -eq $null) -or
                ($retryableStatuses -contains $status -and -not ($status -eq 429 -and -not $hasNumericRetryAfter)) -and
                ($status -ne 401) -and
                ($status -ne 403)

            if (-not $shouldRetry -or $attempt -ge $MaxAttempts) { throw }

            # Backoff with jitter; ensure *integer* milliseconds for Start-Sleep
            $base = [Math]::Min(5000.0, [double]$BaseDelayMs * [Math]::Pow(2.0, ($attempt - 1)))
            $jitter = Get-Random -Minimum 0 -Maximum 200
            $backoff = $base + $jitter
            $delayMs = if ($hasNumericRetryAfter) { [Math]::Max([double]$retryAfterMs, $backoff) } else { $backoff }
            $delayInt = [int]([Math]::Round($delayMs))
            Start-Sleep -Milliseconds $delayInt
        }
    }
}

function Invoke-BeaconApiPost {
    [CmdletBinding()]
    param([Parameter(Mandatory)][string]$Path,[Parameter(Mandatory)][hashtable]$Body)
    $uri = "{0}{1}" -f $script:BeamState.ApiRoot.TrimEnd('/'), $Path
    $headers = @{
        'Accept'       = 'application/json'
        'Content-Type' = 'application/json'
        'User-Agent'   = "ScriptBeacon/1.0 (PowerShell $($PSVersionTable.PSVersion); $([Environment]::OSVersion.VersionString))"
    }
    if ($script:BeamState.WriteSecret) {
        # tolerant to backend variants
        $headers['X-Beacon-Write'] = $script:BeamState.WriteSecret
        $headers['X-Api-Key']      = $script:BeamState.WriteSecret
        $headers['Authorization']  = "ApiKey $($script:BeamState.WriteSecret)"
    }
    $json = $Body | ConvertTo-Json -Compress -Depth 10
    Invoke-WithRetry -Action { Invoke-RestMethod -Method Post -Uri $uri -Headers $headers -Body $json -ErrorAction Stop }
}

function Invoke-S3PutWithHeaders {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)][string]$Url,
        [Parameter(Mandatory)][hashtable]$Headers,
        [Parameter(Mandatory)][string]$InFile
    )
    $hdr=@{}; foreach($k in $Headers.Keys){ $hdr[$k]=[string]$Headers[$k] }
    $old=$global:ProgressPreference
    try {
        $global:ProgressPreference='SilentlyContinue'
        $resp = Invoke-WithRetry -Action { Invoke-WebRequest -UseBasicParsing -ErrorAction Stop -Method Put -Uri $Url -Headers $hdr -InFile $InFile }
        $etag = $null; if ($resp -and $resp.Headers -and $resp.Headers['ETag']) { $etag = $resp.Headers['ETag'].ToString().Trim('"') }
        $etag
    } finally { $global:ProgressPreference=$old }
}

# ------------------------ Core (internal) ------------------------
function Initialize-BeamSession {
    if (-not $script:BeamState.TranscriptPath) {
        $ts   = (Get-Date).ToUniversalTime().ToString('yyyyMMdd-HHmmssfff')
        $uniq = [System.Guid]::NewGuid().ToString('N').Substring(0,8)
        $name = "beam-$ts-$PID-$uniq.log"
        $script:BeamState.TranscriptPath = [IO.Path]::Combine([IO.Path]::GetTempPath(), $name)
    }
    if (-not $script:BeamState.StartedAtUtc) { $script:BeamState.StartedAtUtc = New-IsoNowUtc }

    # Crash guards (process-wide, registered once)
    if ($script:BeamState.AggregatePerProcess) {
        if (-not $script:BeamState.ProcessExitHookRegistered) {
            $null = Ensure-EventSubscription -SourceIdentifier 'ScriptBeacon.ProcessExit' -RegisterAction {
                Register-ObjectEvent -InputObject ([AppDomain]::CurrentDomain) `
                                     -EventName   'ProcessExit' `
                                     -SourceIdentifier 'ScriptBeacon.ProcessExit' `
                                     -Action { try { & (Get-Command Finalize-Beam -ErrorAction Stop) -Reason 'process_exit' } catch {} } | Out-Null
            }
            $script:BeamState.ProcessExitHookRegistered = $true
        }
        if (-not $script:BeamState.UnhandledHookRegistered) {
            $null = Ensure-EventSubscription -SourceIdentifier 'ScriptBeacon.UnhandledException' -RegisterAction {
                Register-ObjectEvent -InputObject ([AppDomain]::CurrentDomain) `
                                     -EventName   'UnhandledException' `
                                     -SourceIdentifier 'ScriptBeacon.UnhandledException' `
                                     -Action { try { $global:LASTEXITCODE = 1; & (Get-Command Finalize-Beam -ErrorAction Stop) -Reason 'unhandled_exception' } catch {} } | Out-Null
            }
            $script:BeamState.UnhandledHookRegistered = $true
        }
    }
}


function Resolve-BeamPath { param([string]$Override,[string[]]$Candidates) if ($Override) { ,$Override } else { $Candidates } }
function Format-Path { param([string]$Template) $Template.Replace('{id}', [System.Uri]::EscapeDataString($script:BeamState.BleepId)) }

function Try-Post-First {
    param([string[]]$Templates,[hashtable]$Body,[ref]$UsedPath)
    foreach ($t in $Templates) {
        $p = Format-Path -Template $t
        try {
            $resp = Invoke-BeaconApiPost -Path $p -Body $Body
            $UsedPath.Value = $p
            return $resp
        } catch {
            $code = $null; try { $code = $_.Exception.Response.StatusCode.Value__ } catch {}
            if ($code -in 404,405,410) { continue }
            throw
        }
    }
    throw "All candidate endpoints unavailable:`n - " + ($Templates -join "`n - ")
}

# Remove transcript banners; return path to clean UTF-8 file
# Preserves epilogue (our appended error block) but prunes transcript noise + duplicates.
function Convert-TranscriptToCleanLog {
    [CmdletBinding()] param([Parameter(Mandatory)][ValidateScript({ Test-Path $_ })][string]$TranscriptPath)

    $text  = [System.IO.File]::ReadAllText($TranscriptPath)
    $lines = $text -split "\r?\n", -1
    $len   = $lines.Length
    $outLines = New-Object System.Collections.Generic.List[string]

    $i = 0
    while ($i -lt $len) {
        # Find transcript header ("PowerShell transcript start")
        $startIdx = -1
        for (; $i -lt $len; $i++) {
            if ($lines[$i] -match '(?i)PowerShell transcript start') { $startIdx = $i; break }
        }
        if ($startIdx -lt 0) { break }

        # Header end (****** line)
        $headerEndStar = -1
        for ($j = $startIdx; $j -lt $len; $j++) {
            if ($lines[$j] -match '^\*{6,}$') { $headerEndStar = $j; break }
        }
        if ($headerEndStar -lt 0) { break }
        $bodyStart = $headerEndStar + 1

        # Footer star BEFORE "PowerShell transcript end"
        $footerStar = -1
        for ($k = $bodyStart; $k -lt ($len - 1); $k++) {
            if ($lines[$k] -match '^\*{6,}$' -and $lines[$k + 1] -match '(?i)PowerShell transcript end') { $footerStar = $k; break }
        }

        if ($footerStar -lt 0) {
            # No standard footer; keep everything until EOF
            for ($p = $bodyStart; $p -lt $len; $p++) { $outLines.Add($lines[$p]) }
            break
        }

        # Copy transcript body (between header/footer banners)
        for ($p = $bodyStart; $p -lt $footerStar; $p++) { $outLines.Add($lines[$p]) }

        # --- Epilogue collection (content appended AFTER transcript end) ---
        # Skip: star line, "PowerShell transcript end", optional "End time:", any star banners/blank lines
        $i = $footerStar + 1
        if ($i -lt $len -and $lines[$i] -match '^\*{6,}$') { $i++ }
        if ($i -lt $len -and $lines[$i] -match '(?i)PowerShell transcript end') { $i++ }
        if ($i -lt $len -and $lines[$i] -match '^(?i)\s*End time\s*:') { $i++ }
        while ($i -lt $len -and ($lines[$i] -match '^\*{6,}$' -or [string]::IsNullOrWhiteSpace($lines[$i]))) { $i++ }

        # Collect epilogue until next transcript header (if any) or EOF
        while ($i -lt $len) {
            if ($lines[$i] -match '^\*{6,}$' -and ($i + 1) -lt $len -and $lines[$i + 1] -match '(?i)PowerShell transcript start') { break }
            $outLines.Add($lines[$i])
            $i++
        }
    }

    if ($outLines.Count -eq 0) {
        # Fallback: if structure not recognized, return file as-is
        foreach ($l in $lines) { $outLines.Add($l) }
    }

    # ---- Post-filter: remove transcript noise & redundancy ----
    $filtered = New-Object System.Collections.Generic.List[string]
    foreach ($l in $outLines) {
        if ([string]::IsNullOrEmpty($l)) { $filtered.Add($l); continue }

        # Drop the single-line transcript banner for terminating errors
        if ($l -match '^\s*PS\s+.*?>\s+TerminatingError\(\):') { continue }

        # Drop leftover transcript banners and time lines
        if ($l -match '^\*{6,}$') { continue }
        if ($l -match '(?i)PowerShell transcript (start|end)') { continue }
        if ($l -match '^(?i)\s*(Start|End)\s*time\s*:') { continue }

        # Strip our own control lines if any slipped in
        if ($l -match '^(📡|──|⏸️)') { continue }

        $filtered.Add($l)
    }

    $cleanPath = [IO.Path]::ChangeExtension($TranscriptPath, ".clean.log")
    $utf8NoBom = New-Object System.Text.UTF8Encoding($false)
    [System.IO.File]::WriteAllText($cleanPath, ($filtered -join [Environment]::NewLine), $utf8NoBom)
    return $cleanPath
}

# Finalize: stop transcript, (optionally) append captured error block, upload
function Finalize-Beam {
    [CmdletBinding()]
    param([string]$Reason = 'manual')

    # Only skip when we’re already finalized AND there isn’t an open transcript.
    if ($script:BeamState.Finalized -and -not $script:BeamState.IsOpen) { return }
    
    # small grace so pending pipeline/host buffers flush
    Start-Sleep -Milliseconds 250

    if ($script:BeamState.IsOpen) {
        try { Stop-Transcript | Out-Null } catch {}
        $script:BeamState.IsOpen   = $false
        $script:BeamState.NestDepth= 0
        Write-Verbose "📡 Beam CLOSED"
    }

    # Append the preformatted error block (host-like), if any
    try {
        if ($script:BeamState.PendingErrorBlock -and (Test-Path $script:BeamState.TranscriptPath)) {
            $utf8NoBom = New-Object System.Text.UTF8Encoding($false)
            [System.IO.File]::AppendAllText($script:BeamState.TranscriptPath, [Environment]::NewLine + $script:BeamState.PendingErrorBlock + [Environment]::NewLine, $utf8NoBom)
        }
    } catch {}
    $script:BeamState.PendingErrorBlock = $null

    if (-not (Test-Path $script:BeamState.TranscriptPath)) {
        Write-Verbose "No transcript found; nothing to upload."
        $script:BeamState.Finalized=$true
        return
    }
    if (-not $script:BeamState.BleepId) {
        Write-Verbose "No Beacon Id configured; skipping upload."
        $script:BeamState.Finalized=$true
        return
    }

    $script:BeamState.EndedAtUtc = New-IsoNowUtc

    # Clean + gzip
    $script:BeamState.CleanPath = Convert-TranscriptToCleanLog -TranscriptPath $script:BeamState.TranscriptPath
    $gz = "$($script:BeamState.CleanPath).gz"
    Compress-Gzip -Path $script:BeamState.CleanPath -Destination $gz
    $script:BeamState.GzipPath = $gz

    # NEVER let size_bytes be 0 (API rejects <=0)
    $bytes = ([IO.FileInfo]$gz).Length
    if ($bytes -lt 1) { $bytes = 1 }
    $sha   = (Get-FileHash -Algorithm SHA256 -Path $gz).Hash.ToLowerInvariant()

    # Local helper: parse HTTP error → friendly guard-rail message (no throw)
    function _ExplainIngestError {
        param([System.Management.Automation.ErrorRecord]$Err)

        $status = $null; $retryAfterSec = $null; $rawText = $null; $json = $null

        try {
            $resp = $Err.Exception.Response
            if ($resp) {
                try { $status = [int]$resp.StatusCode } catch {}
                try {
                    $ra = $resp.Headers['Retry-After']
                    if (-not $ra -and $resp -is [System.Net.HttpWebResponse]) { $ra = $resp.GetResponseHeader('Retry-After') }
                    $tmp = 0
                    if ($ra -and [int]::TryParse([string]$ra, [ref]$tmp)) { $retryAfterSec = $tmp }
                } catch {}
                try {
                    $stream = $resp.GetResponseStream()
                    if ($stream) {
                        $reader = New-Object System.IO.StreamReader($stream)
                        $rawText = $reader.ReadToEnd()
                        $reader.Dispose(); $stream.Dispose()
                    }
                } catch {}
            }
        } catch {}
        if (-not $rawText) {
            try {
                if ($Err.ErrorDetails -and $Err.ErrorDetails.Message) { $rawText = $Err.ErrorDetails.Message }
            } catch {}
        }
        if ($rawText) {
            try { $json = $rawText | ConvertFrom-Json -ErrorAction Stop } catch {}
        }

        $code = if ($json -and $json.PSObject.Properties['error']) { [string]$json.error } else { $null }
        $mode = if ($json -and $json.PSObject.Properties['mode'])  { [string]$json.mode  } else { $null }
        $reason = if ($json -and $json.PSObject.Properties['reason']) { [string]$json.reason } else { $null }
        $bucket = if ($json -and $json.PSObject.Properties['bucket']) { [string]$json.bucket } else { $null }
        $message = if ($json -and $json.PSObject.Properties['message']) { [string]$json.message } else { $null }

        # Map to user-facing explanation.
        # Guards & codes from backend:
        # - Throttle: 429 { error: 'bleep_throttle_exceeded', message } (per-beacon) :contentReference[oaicite:2]{index=2}
        # - Pre-auth IP rate gate: 429 { error: 'preauth_rate_limited', bucket }, adds Retry-After header :contentReference[oaicite:3]{index=3}
        # - IP guard: 403 { error: 'ip_not_allowed', mode: 'allow_all'|'deny_all' } :contentReference[oaicite:4]{index=4}
        # - Write secret: 401 { error: 'missing_write_secret' | 'invalid_write_secret' } or 503 { error:'server_misconfig', details } :contentReference[oaicite:5]{index=5}
        # - Disabled beacon: 403 { error: 'bleep_disabled', reason } :contentReference[oaicite:6]{index=6}
        # - Not found: 404 { error: 'beacon_not_found' (aliased) } :contentReference[oaicite:7]{index=7}
        # - Payload guardrails (logbook): 400 { error: 'payload_too_large'|'invalid_event_payload' } :contentReference[oaicite:8]{index=8}

        $kind = 'unknown'; $friendly = $null; $tip = $null

        switch ($status) {
            429 {
                if ($code -eq 'bleep_throttle_exceeded') {
                    $kind = 'throttle'
                    $friendly = if ($message) { $message } else { 'Beacon throttle exceeded.' }
                    # (Per-beacon throttle does not set Retry-After; we still surface clearly.) :contentReference[oaicite:9]{index=9}
                } elseif ($code -eq 'preauth_rate_limited') {
                    $kind = 'rate_limit'
                    $friendly = 'Ingest pre-auth rate limit hit for your IP.'
                    if ($retryAfterSec) { $tip = "Retry after ~${retryAfterSec}s." }
                } else {
                    $kind = 'rate_limit'
                    $friendly = 'Rate limited.'
                    if ($retryAfterSec) { $tip = "Retry after ~${retryAfterSec}s." }
                }
            }
            401 {
                if ($code -eq 'missing_write_secret') {
                    $kind = 'write_secret'
                    $friendly = 'Beacon requires a write secret but none was provided.'
                    $tip = "Set-BeaconConfig -WriteSecret '<token>' (or set `$env:BEACON_WRITE)."  # :contentReference[oaicite:10]{index=10}
                } elseif ($code -eq 'invalid_write_secret') {
                    $kind = 'write_secret'
                    $friendly = 'Write secret was rejected.'
                    $tip = 'Verify Set-BeaconConfig -WriteSecret / $env:BEACON_WRITE.'         # :contentReference[oaicite:11]{index=11}
                } else {
                    $kind = 'auth'; $friendly = 'Unauthorized (401).'
                }
            }
            403 {
                if ($code -eq 'ip_not_allowed') {
                    $kind = 'ip_guard'
                    $friendly = "IP not allowed by beacon IP Guard (mode=$mode)."
                    $tip = 'Update IP Guard allow/deny lists for this beacon.'                # :contentReference[oaicite:12]{index=12}
                } elseif ($code -eq 'bleep_disabled' -or $code -eq 'beacon_disabled') {
                    $kind = 'disabled'
                    $friendly = 'Beacon is disabled' + ($(if ($reason) { " ($reason)" } else { '' })) + '.'
                } else {
                    $kind = 'forbidden'; $friendly = 'Forbidden (403).'
                }
            }
            404 {
                $kind = 'not_found'
                $friendly = 'Beacon Id not found or not visible to this org.'                 # :contentReference[oaicite:13]{index=13}
            }
            400 {
                if ($code -eq 'payload_too_large') {
                    $kind = 'payload'; $friendly = 'JSON payload exceeds 16 KB.'              # :contentReference[oaicite:14]{index=14}
                } elseif ($code -eq 'invalid_event_payload') {
                    $kind = 'payload'; $friendly = 'Only top‑level primitive JSON values are allowed.' # :contentReference[oaicite:15]{index=15}
                } else {
                    $kind = 'bad_request'; $friendly = 'Bad request (400).'
                }
            }
            503 {
                if ($code -eq 'server_misconfig') {
                    $kind = 'server'
                    $friendly = 'Server misconfiguration blocked ingest.'
                    if ($json -and $json.details) { $tip = "Details: $($json.details)" }     # :contentReference[oaicite:16]{index=16}
                } else {
                    $kind = 'server'; $friendly = 'Service unavailable (503).'
                }
            }
            default {
                if ($status -ge 500 -and $status -lt 600) { $kind='server'; $friendly="Server error ($status)." }
                elseif ($status -ge 400 -and $status -lt 500) { $kind='client'; $friendly="Request error ($status)." }
                else { $kind='unknown'; $friendly='Unexpected error.' }
            }
        }

        [pscustomobject]@{
            Status = $status
            Kind   = $kind
            Message= $friendly
            Tip    = $tip
        }
    }

    # 1) INIT (autodiscovery)
    $initTemplates = Resolve-BeamPath -Override $script:BeamState.InitPathOverride -Candidates $script:InitCandidates
    $initUsed=''; $initRaw = $null
    try {
        $initRaw = Try-Post-First -Templates $initTemplates -Body @{
            size_bytes       = $bytes
            content_type     = $script:BeamState.ContentType
            content_encoding = $script:BeamState.ContentEncoding
            sha256_hex       = $sha
        } -UsedPath ([ref]$initUsed)
    } catch {
        $info = _ExplainIngestError -Err $_
        if ($info -and $info.Message) {
            Write-Warning ("🛡️ Beacon safeguards blocked log upload at INIT: {0}" -f $info.Message)
            if ($info.Tip) { Write-Warning (" Tip: {0}" -f $info.Tip) }
            if ($script:BeamState.CleanPath) {
                Write-Warning (" Local log saved: {0}" -f $script:BeamState.CleanPath)
            }
            $script:BeamState.Finalized=$true
            # Hygiene: unregister crash guards
            try { Unregister-Event -SourceIdentifier 'ScriptBeacon.ProcessExit'        -ErrorAction SilentlyContinue } catch {}
            try { Unregister-Event -SourceIdentifier 'ScriptBeacon.UnhandledException' -ErrorAction SilentlyContinue } catch {}
            $script:BeamState.ProcessExitHookRegistered  = $false
            $script:BeamState.UnhandledHookRegistered    = $false
            return
        }
        throw
    }

    $init = Normalize-InitResponse -Resp $initRaw
    if (-not $init.run_id -or -not $init.upload_url) {
        throw "chunk-init schema unexpected (missing run_id or upload_url). Raw: $($initRaw | ConvertTo-Json -Depth 10)"
    }
    $script:BeamState.RunId     = "$($init.run_id)"
    $script:BeamState.UploadId  = if ($init.upload_id) { "$($init.upload_id)" } else { $script:BeamState.RunId }
    $script:BeamState.ObjectKey = if ($init.object_key) { "$($init.object_key)" } else { $null }

    # 2) PUT (S3 presign)
    $etag = Invoke-S3PutWithHeaders -Url $init.upload_url -Headers $init.required_headers -InFile $gz

    # 3) COMPLETE (autodiscovery; fallback tiny JSON)
    $tags = Merge-BeamTags -Base $script:BeamState.Tags -Extra $null
    $completeBody = @{
        run_id           = $script:BeamState.RunId
        upload_id        = $script:BeamState.UploadId
        session_id       = $script:BeamState.UploadId
        id               = $script:BeamState.RunId
        bytes            = $bytes
        size_bytes       = $bytes
        etag             = $etag
        started_at       = $script:BeamState.StartedAtUtc
        ended_at         = $script:BeamState.EndedAtUtc
        exit_code        = $script:BeamState.ExitCode
        content_type     = $script:BeamState.ContentType
        content_encoding = $script:BeamState.ContentEncoding
        sha256_hex       = $sha
        tags             = $tags
    }
    if ($script:BeamState.ObjectKey) {
        $completeBody['object_key'] = $script:BeamState.ObjectKey
        $completeBody['blob_key']   = $script:BeamState.ObjectKey
        $completeBody['key']        = $script:BeamState.ObjectKey
    }

    $completeTemplates = Resolve-BeamPath -Override $script:BeamState.CompletePathOverride -Candidates $script:CompleteCandidates
    $complete = $null; $completed=$false
    try {
        $completeRaw = Try-Post-First -Templates $completeTemplates -Body $completeBody -UsedPath ([ref]([string]$null))
        $complete    = Normalize-CompleteResponse -Resp $completeRaw -FallbackRunId $script:BeamState.RunId -FallbackUploadId $script:BeamState.UploadId -Bytes $bytes -ETag $etag
        $completed   = $true
    } catch {
        # If COMPLETE failed due to guard rails, surface and stop; only use JSON fallback for "endpoint unavailable" cases.
        $info = _ExplainIngestError -Err $_
        if ($info -and $info.Message) {
            Write-Warning ("🛡️ Beacon safeguards blocked log finalize: {0}" -f $info.Message)
            if ($info.Tip) { Write-Warning (" Tip: {0}" -f $info.Tip) }
            if ($script:BeamState.CleanPath) {
                Write-Warning (" Local log saved: {0}" -f $script:BeamState.CleanPath)
            }
            $script:BeamState.Finalized=$true
            try { Unregister-Event -SourceIdentifier 'ScriptBeacon.ProcessExit'        -ErrorAction SilentlyContinue } catch {}
            try { Unregister-Event -SourceIdentifier 'ScriptBeacon.UnhandledException' -ErrorAction SilentlyContinue } catch {}
            $script:BeamState.ProcessExitHookRegistered  = $false
            $script:BeamState.UnhandledHookRegistered    = $false
            return
        }

        Write-Verbose "chunk-complete unavailable; recording metadata via JSON ingest."
        $ingestTemplates = Resolve-BeamPath -Override $script:BeamState.IngestPathOverride -Candidates $script:IngestCandidates
        $metaBody = @{
            kind             = 'run_complete'
            run_id           = $script:BeamState.RunId
            upload_id        = $script:BeamState.UploadId
            object_key       = $script:BeamState.ObjectKey
            bytes            = $bytes
            etag             = $etag
            sha256_hex       = $sha
            started_at       = $script:BeamState.StartedAtUtc
            ended_at         = $script:BeamState.EndedAtUtc
            exit_code        = $script:BeamState.ExitCode
            content_type     = $script:BeamState.ContentType
            content_encoding = $script:BeamState.ContentEncoding
            tags             = $tags
            note             = 'fallback: chunk-complete not available'
        }
        try {
            $null = Try-Post-First -Templates $ingestTemplates -Body $metaBody -UsedPath ([ref]([string]$null))
            $complete = [pscustomobject]@{ run_id=$script:BeamState.RunId; upload_id=$script:BeamState.UploadId; bytes=$bytes; etag=$etag }
        } catch {
            # If even JSON ingest is blocked by guard rails, surface clearly and exit gracefully.
            $i2 = _ExplainIngestError -Err $_
            if ($i2 -and $i2.Message) {
                Write-Warning ("🛡️ Beacon safeguards blocked metadata record: {0}" -f $i2.Message)
                if ($i2.Tip) { Write-Warning (" Tip: {0}" -f $i2.Tip) }
                if ($script:BeamState.CleanPath) {
                    Write-Warning (" Local log saved: {0}" -f $script:BeamState.CleanPath)
                }
                $script:BeamState.Finalized=$true
                try { Unregister-Event -SourceIdentifier 'ScriptBeacon.ProcessExit'        -ErrorAction SilentlyContinue } catch {}
                try { Unregister-Event -SourceIdentifier 'ScriptBeacon.UnhandledException' -ErrorAction SilentlyContinue } catch {}
                $script:BeamState.ProcessExitHookRegistered  = $false
                $script:BeamState.UnhandledHookRegistered    = $false
                return
            }
            throw
        }
    }

    $script:BeamState.Finalized = $true

    # Hygiene: unregister crash guards (safe if not registered)
    try { Unregister-Event -SourceIdentifier 'ScriptBeacon.ProcessExit'        -ErrorAction SilentlyContinue } catch {}
    try { Unregister-Event -SourceIdentifier 'ScriptBeacon.UnhandledException' -ErrorAction SilentlyContinue } catch {}
    $script:BeamState.ProcessExitHookRegistered  = $false
    $script:BeamState.UnhandledHookRegistered    = $false

    if ($completed) { Write-Verbose "✅ Beam complete. run_id=$($complete.run_id) bytes=$($complete.bytes) etag=$($complete.etag)" }
    else            { Write-Verbose "⚠️ Beam metadata recorded (fallback). run_id=$($complete.run_id) bytes=$($complete.bytes) etag=$($complete.etag)" }
}

# ------------------------ Tagging helpers ------------------------
function Add-BeamTags {
    <#
    .SYNOPSIS
      Append/override tags at any time during a run.
    .EXAMPLE
      Add-BeamTags @{ phase='ingest'; shard='2' }
    .EXAMPLE
      Add-BeamTags 'phase=process','attempt=1'
    .EXAMPLE
      Add-BeamTags 'phase=finalize, test=success'
    #>

    [CmdletBinding()]
    param([Parameter(Mandatory)][object]$Tags)
    if (-not $script:BeamState.Tags) { $script:BeamState.Tags = @{} }
    $script:BeamState.Tags = Merge-BeamTags -Base $script:BeamState.Tags -Extra $Tags
}

# ------------------------ Internal Begin/End (not exported) ------------------------
function _OpenCore {
    [CmdletBinding()]
    param([object]$Tags,[string]$Id,[switch]$NoAggregate)

    # Reset run-scoped state for a fresh top-level invocation
    if (-not $script:BeamState.IsOpen -or $script:BeamState.NestDepth -le 0) {
        $script:BeamState.Finalized          = $false
        $script:BeamState.PendingErrorBlock  = $null
        $script:BeamState.StartedAtUtc       = $null
        $script:BeamState.EndedAtUtc         = $null
        $script:BeamState.ExitCode           = 0
        $script:BeamState.RunId              = $null
        $script:BeamState.UploadId           = $null
        $script:BeamState.ObjectKey          = $null
        $script:BeamState.CleanPath          = $null
        $script:BeamState.GzipPath           = $null
        $script:BeamState.TranscriptPath     = $null
    }

    if ($PSBoundParameters.ContainsKey('Id')) { $script:BeamState.BleepId = $Id }
    if ($NoAggregate.IsPresent) { $script:BeamState.AggregatePerProcess = $false }
    if (-not $script:BeamState.BleepId) { throw "Beacon Id is required. Use Set-BeaconConfig -Id <uuid> or set `$env:BLEEP_ID" }

    Initialize-BeamSession

    if ($Tags) {
        if (-not $script:BeamState.Tags) { $script:BeamState.Tags=@{} }
        $script:BeamState.Tags = Merge-BeamTags -Base $script:BeamState.Tags -Extra $Tags
    }

    if ($script:BeamState.IsOpen) {
        $script:BeamState.NestDepth++
        Write-Verbose "Beam segment started (nest=$($script:BeamState.NestDepth))"
        return
    }

    Write-Verbose "📡 Beam OPEN (aggregate=$($script:BeamState.AggregatePerProcess))"

    $append = (Test-Path $script:BeamState.TranscriptPath)
    try {
        if ($append) { Start-Transcript -Path $script:BeamState.TranscriptPath -Append | Out-Null }
        else         { Start-Transcript -Path $script:BeamState.TranscriptPath | Out-Null }
    } catch {
        throw "Unable to Start-Transcript to $($script:BeamState.TranscriptPath): $($_.Exception.Message)"
    }

    $script:BeamState.IsOpen=$true
    $script:BeamState.NestDepth=1
    if (-not $script:BeamState.StartedAtUtc) { $script:BeamState.StartedAtUtc = New-IsoNowUtc }
}

function _CloseCore {
    [CmdletBinding()] param([object]$Tags,[int]$ExitCode,[switch]$Finalize)

    if ($Tags) {
        if (-not $script:BeamState.Tags){ $script:BeamState.Tags=@{} }
        $script:BeamState.Tags = Merge-BeamTags -Base $script:BeamState.Tags -Extra $Tags
    }
    if ($PSBoundParameters.ContainsKey('ExitCode')) { $script:BeamState.ExitCode = $ExitCode }

    if ($script:BeamState.IsOpen) {
        $script:BeamState.NestDepth=[Math]::Max(0,$script:BeamState.NestDepth-1)
        if ($script:BeamState.NestDepth -le 0) {
            if (-not $Finalize.IsPresent) {
                try { Stop-Transcript | Out-Null } catch {}
                $script:BeamState.IsOpen=$false
                Write-Verbose "📡 Beam CLOSED"
            }
        } else {
            Write-Verbose "Beam segment closed (nest=$($script:BeamState.NestDepth))"
        }
    }

    $shouldFinalize = $Finalize.IsPresent -or (-not $script:BeamState.AggregatePerProcess)
    if ($shouldFinalize) {
        Finalize-Beam -Reason 'explicit'
    } else {
        Write-Verbose "Beam paused. Will upload on process exit or when finalized."
    }
}

# ------------------------ Public single-entry API ------------------------
function Open-Beam {
    <#
    .SYNOPSIS
      Run a scriptblock and capture/upload its transcript (errors included), silently by default.
    .EXAMPLE
      Open-Beam {
        Write-Host "Hello"
        throw "boom"
      } -Tags @{app='demo'} -EndTags 'job=night','test=failed'
    .EXAMPLE
      Open-Beam {
        $script:passed = $true
      } -FinalTags { @{ test = if ($script:passed) {'success'} else {'failure'} } }
    #>

    [CmdletBinding()]
    param(
        [Parameter(Mandatory)][ScriptBlock]$ScriptBlock,
        [object]$Tags,
        [Alias('EndTags','CloseTags')][object]$FinalTags,
        [string]$Id
    )

    # Only pass provided values so we don't overwrite Id/Tags with nulls
    $openParams = @{}
    if ($PSBoundParameters.ContainsKey('Tags')) { $openParams['Tags'] = $Tags }
    if ($PSBoundParameters.ContainsKey('Id'))   { $openParams['Id']   = $Id }

    _OpenCore @openParams
    try {
        & $ScriptBlock
    $__lexv = Get-Variable -Name LASTEXITCODE -Scope Global -ErrorAction SilentlyContinue
    $script:BeamState.ExitCode = if ($__lexv -and $null -ne $__lexv.Value) { [int]$__lexv.Value } else { 0 }
    } catch {
        $script:BeamState.ExitCode = 1
        # Build a host-like block now; append it after we stop the transcript
        try { $script:BeamState.PendingErrorBlock = Format-ErrorRecordBlock -ErrorRecord $_ } catch {}
        throw
    } finally {
        # Evaluate FinalTags if it's a scriptblock; otherwise pass as-is
        $endTags = $null
        if ($PSBoundParameters.ContainsKey('FinalTags')) {
            $endTags = $FinalTags
            if ($endTags -is [scriptblock]) {
                try { $endTags = & $endTags } catch { $endTags = $null }
            }
        }
        $closeParams = @{}
        if ($null -ne $endTags) { $closeParams['Tags'] = $endTags }
        _CloseCore @closeParams -Finalize
    }
}

Export-ModuleMember -Function Set-BeaconConfig, Get-BeamStatus, Open-Beam, Add-BeamTags