Common/Invoke-SafeGraphRequest.ps1

<#
.SYNOPSIS
    Graph request wrapper with pagination and transient-error retry.
.DESCRIPTION
    Drop-in replacement for Invoke-MgGraphRequest on list endpoints. Follows
    @odata.nextLink until the collection is exhausted (bounded by -MaxPages)
    and retries transient Graph failures (429 throttling, 503/504) with
    exponential backoff, honoring the Retry-After header when present.

    Collection responses return a hashtable whose 'value' key holds the merged
    items from every page (other top-level keys are carried over from the first
    page). Non-collection responses (no 'value' property) pass through
    unchanged, so the helper is safe for single-object GETs too.

    Without this wrapper, raw Invoke-MgGraphRequest calls silently truncate at
    the server page size — a tenant with more apps/policies/users than one page
    yields incomplete assessment results (#952).
.PARAMETER Uri
    Graph URI (relative like '/v1.0/applications?$top=999' or absolute).
.PARAMETER Method
    HTTP method. Pagination only applies to GET; POST is supported for parity
    so call sites can migrate uniformly.
.PARAMETER Body
    Optional request body, passed through to Invoke-MgGraphRequest.
.PARAMETER Headers
    Optional Graph headers, preserved across pages and retries.
.PARAMETER ExpectCollection
    Reject missing value arrays on known list endpoints, including the first page.
.PARAMETER FirstPageOnly
    Explicitly bounded evidence sampling or access probes only. Does not establish
    complete collection. Assessment configuration lists must not use this switch.
.PARAMETER OutputFilePath
    Download a Graph report file; preserves the SDK file-response behavior.
.PARAMETER MaxPages
    Safety cap on pages followed (default 100). A terminating error is raised when the
    cap is hit so partial data cannot produce assessment conclusions.
.PARAMETER MaxRetries
    Retries per page for transient errors (default 4; ~2/4/8/16s backoff).
.EXAMPLE
    $response = Invoke-SafeGraphRequest -Uri '/v1.0/applications?$select=id,appId&$top=999'
    $apps = $response.value # complete across all pages
#>

function Invoke-SafeGraphRequest {
    [CmdletBinding()]
    [OutputType([hashtable])]
    param(
        [Parameter(Mandatory)]
        [ValidateNotNullOrEmpty()]
        [string]$Uri,

        [Parameter()]
        [ValidateSet('GET', 'POST')]
        [string]$Method = 'GET',

        [Parameter()]
        [object]$Body,

        [Parameter()]
        [hashtable]$Headers,

        [Parameter()]
        [switch]$ExpectCollection,

        [Parameter()]
        [switch]$FirstPageOnly,

        [Parameter()]
        [string]$OutputFilePath,

        [Parameter()]
        [ValidateRange(1, 1000)]
        [int]$MaxPages = 100,

        [Parameter()]
        [ValidateRange(0, 8)]
        [int]$MaxRetries = 4
    )

    $allValues = [System.Collections.Generic.List[object]]::new()
    $firstPage = $null
    $currentUri = $Uri
    $pageCount = 0

    while ($currentUri) {
        $pageCount++
        if ($pageCount -gt $MaxPages) {
            Write-Warning 'GraphCollectionIncomplete: no complete collection available.'; throw "GraphCollectionIncomplete: page cap ($MaxPages) reached. No partial collection may be used for assessment."
        }

        $attempt = 0
        $response = $null
        while ($true) {
            try {
                $requestParams = @{ Uri = $currentUri; Method = $Method; ErrorAction = 'Stop' }
                if ($null -ne $Body) { $requestParams['Body'] = $Body }
                if ($Headers) { $requestParams['Headers'] = $Headers }
                if ($OutputFilePath) { $requestParams['OutputFilePath'] = $OutputFilePath }
                $response = Invoke-MgGraphRequest @requestParams
                break
            } catch {
                $attempt++
                $delay = Get-GraphRetryDelay -ErrorRecord $_ -Attempt $attempt
                if ($null -eq $delay -or $attempt -gt $MaxRetries) {
                    # Exclude query values and object identifiers from diagnostics.
                    $endpoint = ($Uri -split '\?', 2)[0] -replace '^https?://[^/]+', ''
                    $endpoint = $endpoint -replace '(?i)[0-9a-f]{8}(?:-[0-9a-f]{4}){3}-[0-9a-f]{12}', '{id}'
                    $endpoint = $endpoint -replace '[^/]*(@|%40)[^/]*', '{id}'
                    $statusCode = 0
                    try { $statusCode = [int]$_.Exception.Response.StatusCode } catch { Write-Debug 'No response status property.' }
                    if (-not $statusCode) {
                        try { $statusCode = [int]$_.Exception.ResponseStatusCode } catch { Write-Debug 'No SDK status property.' }
                    }
                    $graphCode = 'unavailable'
                    try {
                        $errorBody = $_.ErrorDetails.Message | ConvertFrom-Json -ErrorAction Stop
                        if ($errorBody.error.code -match '^[a-zA-Z0-9_.-]{1,100}$') { $graphCode = $errorBody.error.code }
                    } catch { Write-Debug 'No structured Graph error code.' }
                    Write-Warning "GraphCollectionIncomplete: $Method $endpoint; page=$pageCount; HTTP=$statusCode; code=$graphCode."
                    throw
                }
                Write-Verbose "Invoke-SafeGraphRequest: transient Graph error (attempt $attempt of $MaxRetries), retrying in ${delay}s: $($_.Exception.Message)"
                Start-Sleep -Seconds $delay
            }
        }

        if ($OutputFilePath) { return }
        if ($null -eq $response) { Write-Warning 'GraphCollectionIncomplete: null response.'; throw 'GraphCollectionIncomplete: null response.' }
        if ($null -eq $firstPage) { $firstPage = $response }

        # Non-collection response: nothing to merge, return as-is.
        $hasValue = if ($response -is [System.Collections.IDictionary]) { $response.Contains('value') }
                    else { $null -ne $response.PSObject.Properties['value'] }
        if (-not $hasValue) {
            if ($pageCount -eq 1 -and -not $ExpectCollection) { return $response }
            Write-Warning 'GraphCollectionIncomplete: no complete collection available.'; throw 'GraphCollectionIncomplete: missing or malformed collection page.'
        }

        if ($response.value -isnot [System.Collections.IList]) { Write-Warning 'GraphCollectionIncomplete: no complete collection available.'; throw 'GraphCollectionIncomplete: invalid value array.' }
        foreach ($item in @($response.value)) { $allValues.Add($item) }

        $currentUri = if ($Method -ne 'GET' -or $FirstPageOnly) { $null } elseif ($response -is [System.Collections.IDictionary]) { $response['@odata.nextLink'] }
                      else { $response.'@odata.nextLink' }
    }

    # Rebuild the familiar response shape: first page's metadata + merged value.
    $result = @{}
    if ($firstPage -is [System.Collections.IDictionary]) {
        foreach ($key in $firstPage.Keys) {
            if ($key -ne 'value' -and $key -ne '@odata.nextLink') { $result[$key] = $firstPage[$key] }
        }
    }
    $result['value'] = $allValues
    return $result
}

<#
.SYNOPSIS
    Computes the retry delay for a transient Graph error, or $null if the
    error is not retryable.
.DESCRIPTION
    Inspects an ErrorRecord from Invoke-MgGraphRequest. Returns a delay in
    seconds for 429/503/504 responses — from the Retry-After header when the
    response exposes one, otherwise exponential backoff (2^attempt, capped at
    60s). Returns $null for non-transient errors so callers rethrow instead of
    retrying permission or request failures.
.PARAMETER ErrorRecord
    The caught ErrorRecord.
.PARAMETER Attempt
    1-based retry attempt number, used for the backoff exponent.
.EXAMPLE
    $delay = Get-GraphRetryDelay -ErrorRecord $_ -Attempt 2
#>

function Get-GraphRetryDelay {
    [CmdletBinding()]
    [OutputType([System.Nullable[int]])]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.ErrorRecord]$ErrorRecord,

        [Parameter(Mandatory)]
        [int]$Attempt
    )

    $statusCode = 0
    $exception = $ErrorRecord.Exception
    if ($exception.PSObject.Properties['Response'] -and $exception.Response) {
        try { $statusCode = [int]$exception.Response.StatusCode } catch { $statusCode = 0 }
    }
    if ($statusCode -eq 0 -and $exception.Message -match 'TooManyRequests|throttl|\b429\b') {
        $statusCode = 429
    } elseif ($statusCode -eq 0 -and $exception.Message -match 'ServiceUnavailable|\b503\b') {
        $statusCode = 503
    } elseif ($statusCode -eq 0 -and $exception.Message -match 'GatewayTimeout|\b504\b') {
        $statusCode = 504
    }

    if ($statusCode -notin @(429, 503, 504)) { return $null }

    # Honor Retry-After when the response surfaces it.
    try {
        $retryAfter = $exception.Response.Headers.RetryAfter
        if ($retryAfter -and $retryAfter.Date) { return [int][math]::Max(1, [math]::Ceiling(($retryAfter.Date - [datetimeoffset]::UtcNow).TotalSeconds)) }
        if ($retryAfter -and $retryAfter.Delta) {
            return [int]([math]::Ceiling($retryAfter.Delta.TotalSeconds) + 1)
        }
    } catch {
        Write-Debug 'Get-GraphRetryDelay: no readable Retry-After header; using exponential backoff.'
    }

    return [int][math]::Min([math]::Pow(2, $Attempt), 60)
}