Add automatic JWT token refresh capability for long running jobs

Tests/TokenRefresh.Tests.ps1

@@ -0,0 +1,108 @@
+BeforeAll {
+    . "$PSScriptRoot/../githound.ps1"
+}
+
+Describe 'Test-GitHubSessionTokenNeedsRefresh' {
+    It 'Returns $false for PAT sessions (no ClientId)' {
+        $session = New-GithubSession -OrganizationName 'TestOrg' -Token 'ghp_faketoken'
+        Test-GitHubSessionTokenNeedsRefresh -Session $session | Should -Be $false
+    }
+
+    It 'Returns $false when token has more than 10 minutes remaining' {
+        $session = New-GithubSession `
+            -OrganizationName 'TestOrg' `
+            -Token 'ghs_faketoken' `
+            -ClientId 'Iv1.abc123' `
+            -InstallationId '12345' `
+            -PrivateKeyPath '/tmp/fake.pem' `
+            -TokenExpiresAt ([System.DateTimeOffset]::UtcNow.AddMinutes(30))
+
+        Test-GitHubSessionTokenNeedsRefresh -Session $session | Should -Be $false
+    }
+
+    It 'Returns $true when token expires within 10 minutes' {
+        $session = New-GithubSession `
+            -OrganizationName 'TestOrg' `
+            -Token 'ghs_faketoken' `
+            -ClientId 'Iv1.abc123' `
+            -InstallationId '12345' `
+            -PrivateKeyPath '/tmp/fake.pem' `
+            -TokenExpiresAt ([System.DateTimeOffset]::UtcNow.AddMinutes(5))
+
+        Test-GitHubSessionTokenNeedsRefresh -Session $session | Should -Be $true
+    }
+
+    It 'Returns $true when token is already expired' {
+        $session = New-GithubSession `
+            -OrganizationName 'TestOrg' `
+            -Token 'ghs_faketoken' `
+            -ClientId 'Iv1.abc123' `
+            -InstallationId '12345' `
+            -PrivateKeyPath '/tmp/fake.pem' `
+            -TokenExpiresAt ([System.DateTimeOffset]::UtcNow.AddMinutes(-5))
+
+        Test-GitHubSessionTokenNeedsRefresh -Session $session | Should -Be $true
+    }
+}
+
+Describe 'Update-GitHubSessionToken' {
+    It 'Throws when called on a PAT session' {
+        $session = New-GithubSession -OrganizationName 'TestOrg' -Token 'ghp_faketoken'
+        { Update-GitHubSessionToken -Session $session } | Should -Throw '*not created with GitHub App JWT credentials*'
+    }
+}
+
+Describe 'New-GitHubJwtSession session properties' {
+    It 'Stores JWT credentials on the session object when mocked' {
+        # Mock the REST call that exchanges JWT for installation token
+        Mock Invoke-GithubRestMethod {
+            return [PSCustomObject]@{
+                token      = 'ghs_mockinstallationtoken'
+                expires_at = ([System.DateTimeOffset]::UtcNow.AddHours(1)).ToString('o')
+            }
+        }
+
+        # Mock Get-Content to return a valid PEM (we won't actually sign since REST is mocked)
+        # We need to mock the RSA signing path, so instead mock the whole JWT exchange
+        # Since New-GitHubJwtSession does RSA signing before the REST call, we need a real key
+        # Generate a throwaway RSA key for testing
+        $rsa = [System.Security.Cryptography.RSA]::Create(2048)
+        $pem = $rsa.ExportRSAPrivateKeyPem()
+        $tempPem = [System.IO.Path]::GetTempFileName()
+        Set-Content -Path $tempPem -Value $pem -NoNewline
+
+        try {
+            $session = New-GitHubJwtSession -OrganizationName 'TestOrg' -ClientId 'Iv1.testclient' -PrivateKeyPath $tempPem -AppId '99999'
+
+            $session.ClientId | Should -Be 'Iv1.testclient'
+            $session.PrivateKeyPath | Should -Be $tempPem
+            $session.InstallationId | Should -Be '99999'
+            $session.TokenExpiresAt | Should -Not -BeNullOrEmpty
+            $session.Headers['Authorization'] | Should -Be 'Bearer ghs_mockinstallationtoken'
+        }
+        finally {
+            Remove-Item $tempPem -ErrorAction SilentlyContinue
+        }
+    }
+}
+
+Describe 'Invoke-GithubRestMethod 401 handling' {
+    It 'Does not attempt token refresh on 401 for PAT sessions' {
+        Mock Invoke-WebRequest {
+            $errorDetails = @{ status = "401"; message = "Bad credentials" } | ConvertTo-Json
+            $exception = [System.Net.Http.HttpRequestException]::new("401 Unauthorized")
+            $errorRecord = [System.Management.Automation.ErrorRecord]::new($exception, 'HttpError', 'ConnectionError', $null)
+            $errorRecord | Add-Member -NotePropertyName 'ErrorDetails' -NotePropertyValue ([System.Management.Automation.ErrorDetails]::new($errorDetails)) -Force
+            throw $errorRecord
+        }
+
+        Mock Update-GitHubSessionToken { }
+
+        $session = New-GithubSession -OrganizationName 'TestOrg' -Token 'ghp_faketoken'
+
+        # Should error (401 with no JWT credentials = unrecoverable)
+        Invoke-GithubRestMethod -Session $session -Path 'test' -ErrorVariable restErrors 2>$null
+
+        Should -Invoke Update-GitHubSessionToken -Times 0
+    }
+}

githound.ps1

@@ -6,7 +6,10 @@ function Get-GitHoundFunctionBundle {
         'Normalize-Null',
         'New-GitHoundNode',
         'New-GitHoundEdge',
+        'New-GithubSession',
         'Invoke-GithubRestMethod',
+        'Test-GitHubSessionTokenNeedsRefresh',
+        'Update-GitHubSessionToken',
         'Get-GitHoundRestErrorInfo',
         'Write-GitHoundRestSkipWarning',
         'Wait-GithubRestRateLimit',
@@ -76,7 +79,10 @@ function New-GithubSession {
 
         [Parameter(Mandatory = $false)]
         [string]
-        $PrivateKeyPath
+        $PrivateKeyPath,
+
+        [Parameter(Mandatory = $false)]
+        $TokenExpiresAt
     )
 
     if($Headers['Accept']) {
@@ -120,6 +126,7 @@ function New-GithubSession {
         PrivateKeyPath = $PrivateKeyPath
         HasPersonalAccessToken = ($null -ne $PatHeaders)
         TargetType = if ($EnterpriseName) { 'Enterprise' } elseif ($OrganizationName) { 'Organization' } else { 'Application' }
+        TokenExpiresAt = $TokenExpiresAt
     }
 }
 
@@ -221,11 +228,76 @@ function New-GitHubJwtSession
         -PatHeaders $patHeaders `
         -ClientId $ClientId `
         -InstallationId $InstallationId `
-        -PrivateKeyPath $PrivateKeyPath
-
+        -PrivateKeyPath $PrivateKeyPath `
+        -TokenExpiresAt ([System.DateTimeOffset]::Parse($result.expires_at))
+
     Write-Output $session
 }
 
+function Test-GitHubSessionTokenNeedsRefresh {
+    param(
+        [Parameter(Position = 0, Mandatory = $true)]
+        [PSTypeName('GitHound.Session')]
+        $Session
+    )
+
+    # Only applicable to JWT-based sessions
+    if (-not $Session.ClientId) {
+        return $false
+    }
+
+    if (-not $Session.TokenExpiresAt) {
+        return $false
+    }
+
+    # Refresh if token expires within the next 10 minutes
+    $refreshThreshold = [System.DateTimeOffset]::UtcNow.AddMinutes(10)
+    return $Session.TokenExpiresAt -le $refreshThreshold
+}
+
+function Update-GitHubSessionToken {
+    param(
+        [Parameter(Position = 0, Mandatory = $true)]
+        [PSTypeName('GitHound.Session')]
+        $Session
+    )
+
+    if (-not $Session.ClientId) {
+        throw "Cannot refresh token: session was not created with GitHub App JWT credentials"
+    }
+
+    Write-Host "Refreshing GitHub App installation token..."
+
+    # Generate a new JWT
+    $header = [Convert]::ToBase64String([System.Text.Encoding]::UTF8.GetBytes((ConvertTo-Json -InputObject @{
+        alg = "RS256"
+        typ = "JWT"
+    }))).TrimEnd('=').Replace('+', '-').Replace('/', '_')
+
+    $payload = [Convert]::ToBase64String([System.Text.Encoding]::UTF8.GetBytes((ConvertTo-Json -InputObject @{
+        iat = [System.DateTimeOffset]::UtcNow.AddSeconds(-10).ToUnixTimeSeconds()
+        exp = [System.DateTimeOffset]::UtcNow.AddMinutes(10).ToUnixTimeSeconds()
+        iss = $Session.ClientId
+    }))).TrimEnd('=').Replace('+', '-').Replace('/', '_')
+
+    $rsa = [System.Security.Cryptography.RSA]::Create()
+    $rsa.ImportFromPem((Get-Content $Session.PrivateKeyPath -Raw))
+
+    $signature = [Convert]::ToBase64String($rsa.SignData([System.Text.Encoding]::UTF8.GetBytes("$header.$payload"), [System.Security.Cryptography.HashAlgorithmName]::SHA256, [System.Security.Cryptography.RSASignaturePadding]::Pkcs1)).TrimEnd('=').Replace('+', '-').Replace('/', '_')
+
+    $jwt = "$header.$payload.$signature"
+
+    # Exchange JWT for a new installation access token using a temporary session
+    $jwtsession = New-GithubSession -Token $jwt
+    $result = Invoke-GithubRestMethod -Session $jwtsession -Path "app/installations/$($Session.InstallationId)/access_tokens" -Method POST
+
+    # Update the existing session in-place
+    $Session.Headers['Authorization'] = "Bearer $($result.token)"
+    $Session.TokenExpiresAt = [System.DateTimeOffset]::Parse($result.expires_at)
+
+    Write-Host "GitHub App installation token refreshed. New expiry: $($Session.TokenExpiresAt)"
+}
+
 function Get-GitHubAppInstallation
 {
     [CmdletBinding()]
@@ -405,9 +477,15 @@ function Invoke-GithubRestMethod {
     $LinkHeader = $Null;
     try {
         do {
+            # Proactively refresh token if nearing expiry
+            if (Test-GitHubSessionTokenNeedsRefresh -Session $Session) {
+                Update-GitHubSessionToken -Session $Session
+            }
+
             $requestSuccessful = $false
             $retryCount = 0
-
+            $tokenRefreshed = $false
+
             while (-not $requestSuccessful -and $retryCount -lt 3) {
                 try {
                     if($LinkHeader) {
@@ -426,14 +504,21 @@ function Invoke-GithubRestMethod {
                         Wait-GithubRestRateLimit -Session $Session
                         $retryCount++
                     }
+                    elseif ($errorInfo.Status -eq "401" -and -not $tokenRefreshed -and $Session.ClientId) {
+                        Write-Warning "Received 401 Unauthorized. Attempting token refresh..."
+                        Update-GitHubSessionToken -Session $Session
+                        $Headers = $Session.Headers
+                        $tokenRefreshed = $true
+                        $retryCount++
+                    }
                     else {
                         throw $_
                     }
                 }
             }
-            
+
             if (-not $requestSuccessful) {
-                throw "Failed after 3 retry attempts due to rate limiting"
+                throw "Failed after 3 retry attempts due to rate limiting or authentication failure"
             }
 
 
@@ -483,6 +568,11 @@ function Invoke-GitHubGraphQL
         $Variables
     )
 
+    # Proactively refresh token if nearing expiry
+    if (Test-GitHubSessionTokenNeedsRefresh -Session $Session) {
+        Update-GitHubSessionToken -Session $Session
+    }
+
     if (-not $Headers) {
         $Headers = $Session.Headers
     }
@@ -501,6 +591,7 @@ function Invoke-GitHubGraphQL
     $requestSuccessful = $false
     $retryCount = 0
     $maxRetries = 5
+    $tokenRefreshed = $false
 
     while (-not $requestSuccessful -and $retryCount -lt $maxRetries) {
         try {
@@ -510,13 +601,17 @@ function Invoke-GitHubGraphQL
         catch {
             $isRateLimit = $false
             $isRetryable = $false
+            $isUnauthorized = $false
             $errorString = "$($_.Exception.Message) $($_.ErrorDetails)"
 
             try {
                 $httpException = $_.ErrorDetails | ConvertFrom-Json
                 if (($httpException.status -eq "403" -and $httpException.message -match "rate limit") -or $httpException.status -eq "429") {
                     $isRateLimit = $true
                 }
+                if ($httpException.status -eq "401") {
+                    $isUnauthorized = $true
+                }
                 if ($httpException.message -match "couldn.t respond.*in time" -or $httpException.message -match "timeout") {
                     $isRetryable = $true
                 }
@@ -526,16 +621,26 @@ function Invoke-GitHubGraphQL
                 if ($errorString -match "rate limit" -or $errorString -match "abuse" -or $errorString -match "secondary" -or $errorString -match "429") {
                     $isRateLimit = $true
                 }
+                if ($errorString -match "401" -or $errorString -match "Bad credentials" -or $errorString -match "Unauthorized") {
+                    $isUnauthorized = $true
+                }
             }
 
             # Catch server errors (502, 503), timeouts, and gateway errors as retryable
-            if (-not $isRateLimit -and -not $isRetryable) {
+            if (-not $isRateLimit -and -not $isRetryable -and -not $isUnauthorized) {
                 if ($errorString -match "502" -or $errorString -match "503" -or $errorString -match "Bad Gateway" -or $errorString -match "couldn.t respond.*in time" -or $errorString -match "timeout") {
                     $isRetryable = $true
                 }
             }
 
-            if ($isRateLimit) {
+            if ($isUnauthorized -and -not $tokenRefreshed -and $Session.ClientId) {
+                Write-Warning "Received 401 Unauthorized on GraphQL call. Attempting token refresh..."
+                Update-GitHubSessionToken -Session $Session
+                $fparams.Headers = $Session.Headers
+                $tokenRefreshed = $true
+                $retryCount++
+            }
+            elseif ($isRateLimit) {
                 Write-Warning "Rate limit hit when doing GraphQL call. Retry $($retryCount + 1)/$maxRetries"
                 Write-Debug $_
                 Wait-GithubGraphQlRateLimit -Session $Session
@@ -554,7 +659,7 @@ function Invoke-GitHubGraphQL
     }
 
     if (-not $requestSuccessful) {
-        throw "Failed after $maxRetries retry attempts due to server errors or rate limiting"
+        throw "Failed after $maxRetries retry attempts due to server errors, rate limiting, or authentication failure"
     }
 
     return $result