Private/Graph/Invoke-GraphBatchRequest.ps1

# Copyright (c) 2026 Sandy Zeng. All rights reserved.
# Source-available. All rights reserved. See LICENSE file.

<#
    Invoke-GraphBatchRequest.ps1 — Sends batched Microsoft Graph API requests (up to 20 per call) with retry logic.
 
    Author: Sandy Zeng
    Project: IntuneDiff
 
    Version History:
    1.0.0 Initial release.
#>


function Invoke-GraphBatchRequest {
    <#
    .SYNOPSIS
        Sends a batch of up to 20 Graph API requests in a single POST /$batch call.
 
    .DESCRIPTION
        Microsoft Graph JSON Batching allows up to 20 requests per batch.
        This function handles:
        - Splitting into chunks of 20 if more are provided
        - Per-item 429 throttling with retry
        - Per-item error detection
        - Token refresh on 401
 
    .PARAMETER Requests
        Array of hashtables, each with: id, method, url
        Example: @{ id = '1'; method = 'GET'; url = '/beta/deviceManagement/...' }
 
    .PARAMETER MaxRetries
        Maximum retries for throttled (429) individual responses. Default 3.
 
    .OUTPUTS
        Hashtable keyed by request id. Each value is the response body (parsed JSON)
        or $null if the request failed.
    #>

    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [object[]]$Requests,

        [int]$MaxRetries = 3
    )

    $batchSize = 20
    $results = @{}

    # Split into chunks of 20
    $chunks = [System.Collections.Generic.List[object[]]]::new()
    for ($i = 0; $i -lt $Requests.Count; $i += $batchSize) {
        $end = [Math]::Min($i + $batchSize, $Requests.Count)
        $chunk = $Requests[$i..($end - 1)]
        $chunks.Add($chunk)
    }

    foreach ($chunk in $chunks) {
        $pendingRequests = @($chunk)
        $retryCount = 0

        while ($pendingRequests.Count -gt 0 -and $retryCount -le $MaxRetries) {
            $batchBody = @{
                requests = @($pendingRequests | ForEach-Object {
                    @{
                        id     = [string]$_.id
                        method = $_.method
                        url    = $_.url
                    }
                })
            }

            Write-IDLog "Graph BATCH: $($pendingRequests.Count) requests (attempt $($retryCount + 1))"

            try {
                $jsonBody = $batchBody | ConvertTo-Json -Depth 10 -Compress
                $response = Invoke-MgGraphRequest -Method POST -Uri 'https://graph.microsoft.com/beta/$batch' `
                    -Body $jsonBody `
                    -Headers @{ 'Content-Type' = 'application/json' } `
                    -OutputType Hashtable
            } catch {
                Write-IDLog "Batch POST failed: $($_.Exception.Message)"
                # If the entire batch fails with 401, try token refresh
                $statusCode = $null
                if ($_.Exception.Response) {
                    $statusCode = [int]$_.Exception.Response.StatusCode
                }
                if ($statusCode -eq 401 -or $_.Exception.Message -match 'Unauthorized') {
                    Write-IDLog 'Batch 401 - refreshing token...'
                    # Use the same refresh logic as Invoke-IntuneDiffRequest
                    $app = $script:MSALApp
                    if ($app) {
                        $accounts = $app.GetAccountsAsync().GetAwaiter().GetResult()
                        $account = $accounts | Select-Object -First 1
                        if ($account) {
                            $scopes = Get-IntuneDiffRequiredScopes
                            $silentBuilder = $app.AcquireTokenSilent($scopes, $account)
                            $authResult = $silentBuilder.ExecuteAsync().GetAwaiter().GetResult()
                            if ($authResult -and $authResult.AccessToken) {
                                Write-MSALCache
                                $secureToken = ConvertTo-SecureString $authResult.AccessToken -AsPlainText -Force
                                try { Disconnect-MgGraph -ErrorAction SilentlyContinue | Out-Null } catch {}
                                try { Connect-MgGraph -AccessToken $secureToken -NoWelcome | Out-Null } catch {
                                    Connect-MgGraph -AccessToken $secureToken | Out-Null
                                }
                                $retryCount++
                                continue
                            }
                        }
                    }
                }
                # If we can't recover, mark all pending as failed
                foreach ($req in $pendingRequests) {
                    $results[[string]$req.id] = $null
                }
                break
            }

            # Process individual responses
            $throttled = [System.Collections.Generic.List[object]]::new()
            $maxRetryAfter = 0

            # Invoke-MgGraphRequest may return IDictionary or PSObject - handle both
            $responses = $null
            if ($response -is [System.Collections.IDictionary]) {
                $responses = $response['responses']
            } elseif ($response.PSObject.Properties['responses']) {
                $responses = $response.responses
            }

            if (-not $responses -or @($responses).Count -eq 0) {
                Write-IDLog "Batch response had no 'responses' array. Marking all as failed."
                foreach ($req in $pendingRequests) {
                    $results[[string]$req.id] = $null
                }
                break
            }

            foreach ($resp in $responses) {
                $id = $null
                $status = 0
                $body = $null

                # Handle both IDictionary and PSObject response items
                if ($resp -is [System.Collections.IDictionary]) {
                    $id = [string]$resp['id']
                    $status = [int]$resp['status']
                    $body = $resp['body']
                } else {
                    $id = [string]$resp.id
                    $status = [int]$resp.status
                    $body = $resp.body
                }

                if ($status -ge 200 -and $status -lt 300) {
                    $results[$id] = $body
                } elseif ($status -eq 429) {
                    # Throttled - queue for retry
                    $retryAfter = 2
                    $headers = if ($resp -is [System.Collections.IDictionary]) { $resp['headers'] } else { $resp.headers }
                    if ($headers) {
                        $ra = if ($headers -is [System.Collections.IDictionary]) { $headers['Retry-After'] } else { $headers.'Retry-After' }
                        if ($ra) { try { $retryAfter = [int]$ra } catch {} }
                    }
                    if ($retryAfter -gt $maxRetryAfter) { $maxRetryAfter = $retryAfter }
                    $original = $pendingRequests | Where-Object { [string]$_.id -eq $id } | Select-Object -First 1
                    if ($original) { $throttled.Add($original) }
                } else {
                    # Failed (404, 403, 500, 503, 504, etc.) - store null, let caller fallback
                    Write-IDLog "Batch item $id failed with status $status"
                    $results[$id] = $null
                }
            }

            if ($throttled.Count -gt 0 -and $retryCount -lt $MaxRetries) {
                Write-IDLog "Batch: $($throttled.Count) items throttled, retrying after ${maxRetryAfter}s..."
                Start-Sleep -Seconds ([Math]::Max($maxRetryAfter, 1))
                $pendingRequests = @($throttled.ToArray())
                $retryCount++
            } else {
                # Mark any remaining throttled as failed
                foreach ($req in $throttled) {
                    $results[[string]$req.id] = $null
                }
                break
            }
        }
    }

    return $results
}