Script to assign copilot budgets based on EntraID Group Membership

assign-copilot-budgets.ps1

$ErrorActionPreference = 'Stop'

if ($MyInvocation.InvocationName -ne '.') {
    Import-Module Microsoft.Graph.Authentication
    Import-Module Microsoft.Graph.Applications
    Import-Module Microsoft.Graph.Users
    Import-Module Microsoft.Graph.Groups
}

function Get-EntraIdUser {
    param(
        [Parameter(ValueFromPipeline = $true)]
        [string[]]$mails,

        $enterpriseUser
    )

    if ($mails.Count -gt 0) {
        foreach ($mail in $mails) {
            $entraUser = $entraCache[$mail.ToLowerInvariant()]
            if ($entraUser) {
                return $entraUser
            }
        }
    }
    else {
        $mails = @() 
        if ($enterpriseUser.github_com_saml_name_id) {
            $mails += $enterpriseUser.github_com_saml_name_id
        }
        if ($enterpriseUser.github_com_login -match "@") {
            $mails += $enterpriseUser.github_com_login
        }

        $mails += $enterpriseUser.github_com_verified_domain_emails
        return Get-EntraIdUser -mails ($mails | ? { $_ -ne $null } | select-object -Unique -CaseInsensitive)
    }
}

function Invoke-Gh {
    <#
    .Synopsis
    Wrapper function that deals with Powershell's peculiar error output when gh uses the error stream.
    #>
    [CmdletBinding()]
    param(
        [Parameter(Mandatory = $false)]
        [switch]$slurp,

        [Parameter(Mandatory = $false)]
        [switch]$fromJson,

        [parameter(ValueFromRemainingArguments = $true)]
        [string[]]$Arguments
    )

    $output = & {
        [CmdletBinding()]
        param(
            [parameter(ValueFromRemainingArguments = $true)]
            [string[]]$InnerArgs
        )
        $sleep = 0
        do {
            if ($sleep -gt 0) {
                Start-Sleep -Seconds $sleep 
            }
            
            $allOutput = & gh @InnerArgs 2>&1
            $stderr = $allOutput | ? { $_ -is [System.Management.Automation.ErrorRecord] }
            $output = $allOutput | ? { $_ -isnot [System.Management.Automation.ErrorRecord] }

            $sleep += 5
        }
        while (
            $stderr -like "*GraphQL: was submitted too quickly*" -or
            $stderr -like "*API rate limit exceeded*"
        )

        if ($LASTEXITCODE -ne 0) {
            throw "Error executing gh command: $InnerArgs `n" + $stderr
        }

        if ($slurp) {
            $output = $output | & jq -s
        }

        if ($fromJson) {
            $output = $output | ConvertFrom-Json
        }
        return $output
    } -ErrorAction SilentlyContinue -ErrorVariable fail -InnerArgs:$Arguments

    if ($fail) {
        throw $fail.Exception
    }
        
    return $output
}

function Get-BudgetProductSku {
    param(
        [datetime]$AsOf = [datetime]::UtcNow
    )
    return "ai_credits"
}

# Determine the universal budget amount based on the current date.
# June 1 - August 31 2026: $70
# September 1 2026 onwards: $39
function Get-UniversalBudgetAmount {
    param(
        [datetime]$AsOf = [datetime]::UtcNow
    )
    $sept1 = [datetime]::new(2026, 9, 1, 0, 0, 0, [System.DateTimeKind]::Utc)

    if ($AsOf -lt $sept1) {
        return [decimal]70.00
    }
    else {
        return [decimal]39.00
    }
}

# Get Entra ID groups matching the copilot-budget-* naming pattern.
# Returns a hashtable mapping Entra user ID -> target budget amount (highest wins).
# Only groups owned by the designated admin account are trusted.
$script:RequiredGroupOwner = "github-admin@xebia.com"

function Get-CopilotBudgetGroupAssignments {
    Write-Host "Reading Entra ID groups matching 'copilot-budget-*'..."
    $budgetGroups = Get-MgGroup -Filter "startsWith(displayName, 'copilot-budget-')" -All -Property @("Id", "DisplayName")

    $userBudgets = @{}
    $userGroupMemberships = @{}

    foreach ($group in $budgetGroups) {
        # Extract the budget value from the group name (e.g. copilot-budget-100 -> 100)
        if ($group.DisplayName -match '^copilot-budget-(\d+)$') {
            $budgetValue = [decimal]$Matches[1]
        }
        else {
            Write-Warning "Group '$($group.DisplayName)' does not match expected pattern 'copilot-budget-XXXX'. Skipping."
            continue
        }

        # Verify the group is owned by the designated admin to prevent abuse
        $owners = Get-MgGroupOwner -GroupId $group.Id -All
        $ownerUpns = $owners | ForEach-Object {
            $_.AdditionalProperties['userPrincipalName']
        } | Where-Object { $_ -ne $null }

        if ($script:RequiredGroupOwner -notin $ownerUpns) {
            Write-Warning "Group '$($group.DisplayName)' is not owned by $($script:RequiredGroupOwner). Skipping."
            continue
        }

        Write-Host "  Group '$($group.DisplayName)' -> budget `$$budgetValue"
        $members = Get-MgGroupMember -GroupId $group.Id -All

        foreach ($member in $members) {
            $userId = $member.Id
            if (-not $userGroupMemberships.ContainsKey($userId)) { $userGroupMemberships[$userId] = @() }
            $userGroupMemberships[$userId] += $group.DisplayName

            if (-not $userBudgets.ContainsKey($userId) -or $userBudgets[$userId] -lt $budgetValue) {
                $userBudgets[$userId] = $budgetValue
            }
        }
    }

    # Warn about users in multiple budget groups
    foreach ($userId in $userGroupMemberships.Keys) {
        if ($userGroupMemberships[$userId].Count -gt 1) {
            $groups = $userGroupMemberships[$userId] -join ', '
            Write-Warning "User '$userId' is a member of multiple copilot-budget groups: $groups. Applying highest budget: `$$($userBudgets[$userId])."
        }
    }

    return $userBudgets
}

# Get existing Copilot budgets from the GitHub enterprise billing API.
function Get-ExistingBudgets {
    param(
        [string]$Enterprise,
        [string]$ProductSku
    )

    $allBudgets = @()
    $page = 1
    do {
        $response = invoke-gh -fromJson -- api "/enterprises/$Enterprise/settings/billing/budgets?page=$page"
        $allBudgets += $response.budgets
        $page++
    } while ($response.has_next_page)

    return $allBudgets
}

# Create a new individual per-user budget.
function New-IndividualBudget {
    param(
        [string]$Enterprise,
        [string]$GitHubLogin,
        [decimal]$Amount,
        [string]$ProductSku
    )

    $body = @{
        budget_amount         = $Amount
        budget_scope          = "user"
        user                  = $GitHubLogin
        budget_type           = "BundlePricing"
        budget_product_sku    = $ProductSku
        prevent_further_usage = $true
        budget_alerting       = @{
            will_alert       = $false
            alert_recipients = @()
        }
    }

    Write-Output "Creating individual budget of `$$Amount for user $GitHubLogin"
    $_ = ($body | ConvertTo-Json -Depth 5) | gh api --method POST /enterprises/$Enterprise/settings/billing/budgets --input -
}

# Update an existing budget amount.
function Update-Budget {
    param(
        [string]$Enterprise,
        [string]$BudgetId,
        [decimal]$Amount
    )

    $body = @{
        budget_amount         = $Amount
        prevent_further_usage = $true
    }

    Write-Output "Updating budget $BudgetId to `$$Amount"
    $_ = ($body | ConvertTo-Json -Depth 5) | gh api --method PATCH /enterprises/$Enterprise/settings/billing/budgets/$BudgetId --input -
}

# Delete a budget.
function Remove-Budget {
    param(
        [string]$Enterprise,
        [string]$BudgetId
    )

    Write-Output "Deleting budget $BudgetId"
    $_ = gh api --method DELETE /enterprises/$Enterprise/settings/billing/budgets/$BudgetId
}

# Ensure the universal budget exists and has the correct amount.
function Update-UniversalBudget {
    param(
        [string]$Enterprise,
        [string]$ProductSku,
        [decimal]$TargetAmount,
        $ExistingBudgets
    )

    $universalBudget = $ExistingBudgets | Where-Object { $_.budget_scope -eq "multi_user_customer" -and $_.budget_product_sku -eq $ProductSku }

    # Remove any stale universal budget with a different SKU (e.g. premium_requests -> ai_credits cutover)
    $staleUniversalBudgets = $ExistingBudgets | Where-Object { $_.budget_scope -eq "multi_user_customer" -and $_.budget_product_sku -ne $ProductSku }
    foreach ($stale in $staleUniversalBudgets) {
        Write-Output "Removing stale universal budget with SKU '$($stale.budget_product_sku)' (superseded by '$ProductSku')"
        Remove-Budget -Enterprise $Enterprise -BudgetId $stale.id
    }

    if ($universalBudget) {
        if ($universalBudget.budget_amount -ne $TargetAmount) {
            Write-Output "Updating universal budget from `$$($universalBudget.budget_amount) to `$$TargetAmount"
            Update-Budget -Enterprise $Enterprise -BudgetId $universalBudget.id -Amount $TargetAmount
        }
        else {
            Write-Output "Universal budget already set to `$$TargetAmount"
        }
    }
    else {
        Write-Output "Creating universal budget of `$$TargetAmount"
        $body = @{
            budget_amount         = $TargetAmount
            budget_scope          = "multi_user_customer"
            budget_type           = "BundlePricing"
            budget_product_sku    = $ProductSku
            prevent_further_usage = $true
            budget_alerting       = @{
                will_alert       = $false
                alert_recipients = @()
            }
        }
        $_ = ($body | ConvertTo-Json -Depth 5) | gh api --method POST /enterprises/$Enterprise/settings/billing/budgets --input -
    }
}

# Reconcile individual per-user budgets based on Entra ID group membership.
function Update-IndividualBudgets {
    param(
        [string]$Enterprise,
        [string]$ProductSku,
        [hashtable]$TargetUserBudgets,
        $ExistingBudgets,
        $EnterpriseUsers,
        [datetime]$AsOf = [datetime]::UtcNow
    )

    # Get existing individual budgets (scope = "user")
    $existingIndividualBudgets = $ExistingBudgets | Where-Object { $_.budget_scope -eq "user" }

    # Build a lookup of GitHub login -> target budget amount using the Entra cache
    $loginBudgets = @{}
    foreach ($enterpriseUser in $EnterpriseUsers) {
        $entraUser = Get-EntraIdUser -enterpriseUser $enterpriseUser
        if ($entraUser -and $TargetUserBudgets.ContainsKey($entraUser.Id)) {
            $loginBudgets[$enterpriseUser.github_com_login] = $TargetUserBudgets[$entraUser.Id]
        }
    }

    Write-Output "Target individual budgets: $($loginBudgets.Count) users"

    # Create or update budgets for users who should have one
    foreach ($login in $loginBudgets.Keys) {
        $targetAmount = $loginBudgets[$login]
        $existingBudget = $existingIndividualBudgets | Where-Object { $_.budget_entity_name -eq $login }

        if ($existingBudget) {
            if ($existingBudget.budget_product_sku -ne $ProductSku) {
                # The API does not allow changing budget_product_sku in place; remove and recreate.
                Write-Output "Recreating budget for $login (SKU change: '$($existingBudget.budget_product_sku)' -> '$ProductSku')"
                Remove-Budget -Enterprise $Enterprise -BudgetId $existingBudget.id
                New-IndividualBudget -Enterprise $Enterprise -GitHubLogin $login -Amount $targetAmount -ProductSku $ProductSku
            }
            elseif ($existingBudget.budget_amount -ne $targetAmount) {
                Update-Budget -Enterprise $Enterprise -BudgetId $existingBudget.id -Amount $targetAmount
            }
        }
        else {
            New-IndividualBudget -Enterprise $Enterprise -GitHubLogin $login -Amount $targetAmount -ProductSku $ProductSku
        }
    }

    # Remove budgets for users who are no longer in any copilot-budget group
    foreach ($existingBudget in $existingIndividualBudgets) {
        if (-not $loginBudgets.ContainsKey($existingBudget.budget_entity_name)) {
            Write-Output "Removing individual budget for user $($existingBudget.budget_entity_name) (no longer in any copilot-budget group)"
            Remove-Budget -Enterprise $Enterprise -BudgetId $existingBudget.id
        }
    }
}

# --- Main execution (skipped when dot-sourced from Pester tests) ---
$isPesterContext = $MyInvocation.InvocationName -eq '.' -and
                  (Get-PSCallStack | Where-Object { $_.ScriptName -like '*.Tests.ps1' })
if ($isPesterContext) { return }

$enterprise = "xebia"

$accessToken = az account get-access-token --resource https://graph.microsoft.com --query accessToken --output tsv
write-host "::add-mask::$accessToken"
$accessToken = $accessToken | ConvertTo-SecureString -AsPlainText -Force

Connect-MgGraph -AccessToken $accessToken -NoWelcome

# Build Entra ID user cache for mapping GitHub users -> Entra users
Write-Host "Reading Entra ID Users"
# Get-MgUser -All can fail with Directory_ExpiredPageToken (HTTP 400) when pagination
# takes too long and the server-side page token expires. Retry from scratch with
# exponential back-off; each attempt uses a fresh token so there is no stale-token risk.
$entraIdUsers = $null
$maxAttempts = 5
for ($attempt = 1; $attempt -le $maxAttempts; $attempt++) {
    try {
        $entraIdUsers = Get-MgUser -All -ConsistencyLevel eventual -PageSize 999 -Property @("Id", "CompanyName", "DisplayName", "Mail", "UserPrincipalName", "AccountEnabled", "ProxyAddresses") -Filter "UserType eq 'Member'"
        break
    }
    catch {
        $isExpiredToken = $_.Exception.Message -like "*Directory_ExpiredPageToken*" -or
                          $_.Exception.Message -like "*ExpiredPageToken*"
        if (-not $isExpiredToken -or $attempt -eq $maxAttempts) {
            throw
        }
        $backoff = [math]::Pow(2, $attempt) * 5   # 10s, 20s, 40s, 80s
        Write-Warning "Directory_ExpiredPageToken on attempt $attempt/$maxAttempts. Retrying in $backoff seconds..."
        Start-Sleep -Seconds $backoff
    }
}

$entraCache = @{}
foreach ($entraUser in $entraIdUsers) {
    if ($entraUser.Mail) { $entraCache[$entraUser.Mail.ToLowerInvariant()] = $entraUser }
    $entraCache[$entraUser.UserPrincipalName.ToLowerInvariant()] = $entraUser

    foreach ($proxyAddress in $entraUser.ProxyAddresses) {
        $entraCache[$proxyAddress.SubString(5).ToLowerInvariant()] = $entraUser
    }
}

# Get copilot budget group assignments from Entra ID
$targetUserBudgets = Get-CopilotBudgetGroupAssignments

# Get enterprise users from GitHub
Write-Host "Reading GitHub Enterprise Users"
$enterpriseUsers = invoke-gh -- api https://api.github.com/enterprises/$enterprise/consumed-licenses --jq '.users[]' --paginate | ConvertFrom-Json

# Determine product SKU and universal budget amount
$productSku = Get-BudgetProductSku
$universalAmount = Get-UniversalBudgetAmount
Write-Host "Product SKU: $productSku, Universal budget: `$$universalAmount"

# Get existing budgets
$existingBudgets = Get-ExistingBudgets -Enterprise $enterprise -ProductSku $productSku

# Update the universal budget
Update-UniversalBudget -Enterprise $enterprise -ProductSku $productSku -TargetAmount $universalAmount -ExistingBudgets $existingBudgets

# Update individual per-user budgets
Update-IndividualBudgets -Enterprise $enterprise -ProductSku $productSku -TargetUserBudgets $targetUserBudgets -ExistingBudgets $existingBudgets -EnterpriseUsers $enterpriseUsers