Skip to content

Commit 245349b

Browse files
Add error bail-out thresholds to Invoke-AITool
Introduces MaxErrors and MaxTokenErrors parameters to limit general and token/credit-related errors during batch processing. The function now tracks error counts and halts further processing when thresholds are reached, preventing excessive API calls on repeated failures. Runspaces are cleaned up after bail-out, and remaining batches are skipped if limits are exceeded.
1 parent 4e1b13b commit 245349b

File tree

1 file changed

+107
-2
lines changed

1 file changed

+107
-2
lines changed

public/Invoke-AITool.ps1

Lines changed: 107 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,17 @@ function Invoke-AITool {
140140
Example with multiple directories: -ContextFilterBase "C:\primary", "C:\fallback"
141141
Searches C:\primary first, then C:\fallback, then the source file's directory.
142142
143+
.PARAMETER MaxErrors
144+
Maximum number of general errors before bailing out of batch processing. Default is 10.
145+
When this threshold is reached, remaining files/batches are skipped to avoid wasting
146+
API calls on a failing operation. Useful when processing many files in parallel.
147+
148+
.PARAMETER MaxTokenErrors
149+
Maximum number of token/credit-related errors before bailing out. Default is 3.
150+
Token errors are detected by patterns like "token", "credits", "exhausted", "quota",
151+
"insufficient", "billing", "payment". A lower threshold is used because these errors
152+
typically indicate account-wide issues that won't resolve by retrying other files.
153+
143154
.EXAMPLE
144155
Get-ChildItem *.Tests.ps1 | Invoke-AITool -Prompt "Refactor from Pester v4 to v5"
145156
@@ -320,7 +331,13 @@ function Invoke-AITool {
320331
[Parameter()]
321332
[scriptblock]$ContextFilter,
322333
[Parameter()]
323-
[string[]]$ContextFilterBase
334+
[string[]]$ContextFilterBase,
335+
[Parameter()]
336+
[ValidateRange(1, 1000)]
337+
[int]$MaxErrors = 10,
338+
[Parameter()]
339+
[ValidateRange(1, 100)]
340+
[int]$MaxTokenErrors = 3
324341
)
325342

326343
begin {
@@ -585,6 +602,11 @@ function Invoke-AITool {
585602

586603
$filesToProcess = @()
587604
$script:totalInputFiles = 0
605+
606+
# Error tracking for bail-out feature
607+
$script:errorCount = 0
608+
$script:tokenErrorCount = 0
609+
$script:bailedOut = $false
588610
}
589611

590612
process {
@@ -1242,7 +1264,7 @@ function Invoke-AITool {
12421264
# Poll runspaces and output results as they complete (streaming, not buffering)
12431265
$completedBatchCount = 0
12441266
$totalBatches = $batches.Count
1245-
while ($runspaces.Count -gt 0) {
1267+
while ($runspaces.Count -gt 0 -and -not $script:bailedOut) {
12461268
foreach ($runspace in @($runspaces)) {
12471269
if ($runspace.Status.IsCompleted) {
12481270
try {
@@ -1274,12 +1296,54 @@ function Invoke-AITool {
12741296
if ($r.Duration) {
12751297
$null = $allDurations.Add($r.Duration.TotalSeconds)
12761298
}
1299+
# Check for errors and track for bail-out
1300+
if ($r.Success -eq $false) {
1301+
$resultText = $r.Result | Out-String
1302+
if ($resultText -match '(?i)(token|credits?|exhausted|quota|usage.?limit|insufficient|billing|payment)') {
1303+
$script:tokenErrorCount++
1304+
Write-PSFMessage -Level Warning -Message "Token/credit error detected ($script:tokenErrorCount of $MaxTokenErrors max)"
1305+
if ($script:tokenErrorCount -ge $MaxTokenErrors) {
1306+
$script:bailedOut = $true
1307+
Write-PSFMessage -Level Error -Message "BAILING OUT: Reached $MaxTokenErrors token/credit errors. Stopping all processing."
1308+
Write-Warning "BAILING OUT: Reached $MaxTokenErrors token/credit errors. Remaining files will not be processed."
1309+
}
1310+
} else {
1311+
$script:errorCount++
1312+
Write-PSFMessage -Level Warning -Message "Error detected ($script:errorCount of $MaxErrors max)"
1313+
if ($script:errorCount -ge $MaxErrors) {
1314+
$script:bailedOut = $true
1315+
Write-PSFMessage -Level Error -Message "BAILING OUT: Reached $MaxErrors errors. Stopping all processing."
1316+
Write-Warning "BAILING OUT: Reached $MaxErrors errors. Remaining files will not be processed."
1317+
}
1318+
}
1319+
}
12771320
$r
12781321
}
12791322
} else {
12801323
if ($result.Duration) {
12811324
$null = $allDurations.Add($result.Duration.TotalSeconds)
12821325
}
1326+
# Check for errors and track for bail-out
1327+
if ($result.Success -eq $false) {
1328+
$resultText = $result.Result | Out-String
1329+
if ($resultText -match '(?i)(token|credits?|exhausted|quota|usage.?limit|insufficient|billing|payment)') {
1330+
$script:tokenErrorCount++
1331+
Write-PSFMessage -Level Warning -Message "Token/credit error detected ($script:tokenErrorCount of $MaxTokenErrors max)"
1332+
if ($script:tokenErrorCount -ge $MaxTokenErrors) {
1333+
$script:bailedOut = $true
1334+
Write-PSFMessage -Level Error -Message "BAILING OUT: Reached $MaxTokenErrors token/credit errors. Stopping all processing."
1335+
Write-Warning "BAILING OUT: Reached $MaxTokenErrors token/credit errors. Remaining files will not be processed."
1336+
}
1337+
} else {
1338+
$script:errorCount++
1339+
Write-PSFMessage -Level Warning -Message "Error detected ($script:errorCount of $MaxErrors max)"
1340+
if ($script:errorCount -ge $MaxErrors) {
1341+
$script:bailedOut = $true
1342+
Write-PSFMessage -Level Error -Message "BAILING OUT: Reached $MaxErrors errors. Stopping all processing."
1343+
Write-Warning "BAILING OUT: Reached $MaxErrors errors. Remaining files will not be processed."
1344+
}
1345+
}
1346+
}
12831347
$result
12841348
}
12851349
} else {
@@ -1321,6 +1385,20 @@ function Invoke-AITool {
13211385
}
13221386
}
13231387

1388+
# If we bailed out, clean up remaining runspaces
1389+
if ($script:bailedOut -and $runspaces.Count -gt 0) {
1390+
Write-PSFMessage -Level Warning -Message "Cleaning up $($runspaces.Count) remaining runspace(s) after bail-out"
1391+
foreach ($runspace in $runspaces) {
1392+
try {
1393+
$runspace.Pipe.Stop()
1394+
$runspace.Pipe.Dispose()
1395+
} catch {
1396+
Write-PSFMessage -Level Debug -Message "Error disposing runspace: $_"
1397+
}
1398+
}
1399+
$runspaces = @()
1400+
}
1401+
13241402
Write-PSFMessage -Level Verbose -Message "All parallel processing complete"
13251403
Write-Progress -Activity $processingActivity -Completed
13261404

@@ -1342,6 +1420,11 @@ function Invoke-AITool {
13421420

13431421
$batchIndex = 0
13441422
foreach ($batch in $batches) {
1423+
# Check for bail-out before processing each batch
1424+
if ($script:bailedOut) {
1425+
Write-PSFMessage -Level Warning -Message "Skipping remaining batches due to bail-out"
1426+
break
1427+
}
13451428
$batchIndex++
13461429

13471430
# Filter out modified files from this batch if -SkipModified is enabled
@@ -2083,6 +2166,28 @@ function Invoke-AITool {
20832166
}
20842167
}
20852168

2169+
# Check for errors and track for bail-out (sequential mode)
2170+
if ($toolExitCode -ne 0) {
2171+
$resultText = if ($capturedOutput) { $capturedOutput | Out-String } else { '' }
2172+
if ($resultText -match '(?i)(token|credits?|exhausted|quota|usage.?limit|insufficient|billing|payment)') {
2173+
$script:tokenErrorCount++
2174+
Write-PSFMessage -Level Warning -Message "Token/credit error detected ($script:tokenErrorCount of $MaxTokenErrors max)"
2175+
if ($script:tokenErrorCount -ge $MaxTokenErrors) {
2176+
$script:bailedOut = $true
2177+
Write-PSFMessage -Level Error -Message "BAILING OUT: Reached $MaxTokenErrors token/credit errors. Stopping all processing."
2178+
Write-Warning "BAILING OUT: Reached $MaxTokenErrors token/credit errors. Remaining files will not be processed."
2179+
}
2180+
} else {
2181+
$script:errorCount++
2182+
Write-PSFMessage -Level Warning -Message "Error detected ($script:errorCount of $MaxErrors max)"
2183+
if ($script:errorCount -ge $MaxErrors) {
2184+
$script:bailedOut = $true
2185+
Write-PSFMessage -Level Error -Message "BAILING OUT: Reached $MaxErrors errors. Stopping all processing."
2186+
Write-Warning "BAILING OUT: Reached $MaxErrors errors. Remaining files will not be processed."
2187+
}
2188+
}
2189+
}
2190+
20862191
# Apply delay after processing each batch (if not the last batch)
20872192
if ($DelaySeconds -gt 0 -and $batchIndex -lt $batches.Count) {
20882193
Write-PSFMessage -Level Verbose -Message "Waiting $DelaySeconds seconds before processing next batch..."

0 commit comments

Comments
 (0)