Skip to content

Commit ee508da

Browse files
notdpnotdp
committed
fix(api): add retry mechanism for usage fetch to handle SSL failures
1 parent 159b977 commit ee508da

4 files changed

Lines changed: 276 additions & 201 deletions

File tree

‎bin/dk‎

Lines changed: 30 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -40,8 +40,9 @@ load_config() {
4040
load_config
4141

4242
# 可调优的 curl 限制，避免长时间等待慢/无效请求
43-
DKM_CURL_MAX_TIME="${DKM_CURL_MAX_TIME:-4}" # 单请求超时（秒）
44-
DKM_CURL_CONNECT_TIMEOUT="${DKM_CURL_CONNECT_TIMEOUT:-2}" # 连接超时（秒）
43+
DKM_CURL_MAX_TIME="${DKM_CURL_MAX_TIME:-8}" # 单请求超时（秒）
44+
DKM_CURL_CONNECT_TIMEOUT="${DKM_CURL_CONNECT_TIMEOUT:-5}" # 连接超时（秒）
45+
DKM_CURL_RETRIES="${DKM_CURL_RETRIES:-3}" # 重试次数
4546
# 并发与缓存
4647
DKM_LIST_AUTO_MAX="${DKM_LIST_AUTO_MAX:-6}" # 自动模式的最大并发上限
4748
DKM_LIST_JOBS="${DKM_LIST_JOBS:-0}" # 0 表示自动 min(DKM_LIST_AUTO_MAX, key数)
@@ -252,23 +253,35 @@ copy_to_clipboard() {
252253

253254
fetch_usage() {
254255
local key="$1"
255-
local resp http_status body
256-
# 将 HTTP 状态码与 body 分开，非 200 直接标记为不可用（避免继续使用无效 key）
257-
if ! resp=$(curl -sS --retry 0 \
258-
--connect-timeout "$DKM_CURL_CONNECT_TIMEOUT" --max-time "$DKM_CURL_MAX_TIME" \
259-
-w '\n%{http_code}' -X GET 'https://app.factory.ai/api/organization/members/chat-usage' \
260-
-H "Authorization: Bearer ${key}" \
261-
-H 'User-Agent: Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36'); then
262-
printf "BALANCE=0\nBALANCE_NUM=0\nTOTAL=0\nUSED=0\nEXPIRES=?\nEXPIRES_FULL=?\nEXPIRES_TS=\nRAW=curl_error\n"
263-
return
264-
fi
265-
266-
http_status=${resp##*$'\n'}
267-
body=${resp%$'\n'*}
268-
[[ "$http_status" =~ ^[0-9]{3}$ ]] || http_status=000
256+
local resp http_status body attempt
257+
258+
for (( attempt=1; attempt<=DKM_CURL_RETRIES; attempt++ )); do
259+
# 将 HTTP 状态码与 body 分开，非 200 直接标记为不可用（避免继续使用无效 key）
260+
if resp=$(curl -sS --retry 0 \
261+
--connect-timeout "$DKM_CURL_CONNECT_TIMEOUT" --max-time "$DKM_CURL_MAX_TIME" \
262+
-w '\n%{http_code}' -X GET 'https://app.factory.ai/api/organization/members/chat-usage' \
263+
-H "Authorization: Bearer ${key}" \
264+
-H 'User-Agent: Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36' 2>/dev/null); then
265+
266+
http_status=${resp##*$'\n'}
267+
body=${resp%$'\n'*}
268+
[[ "$http_status" =~ ^[0-9]{3}$ ]] || http_status=000
269+
270+
if [ "$http_status" = "200" ]; then
271+
break
272+
elif [[ "$http_status" =~ ^4[0-9]{2}$ ]]; then
273+
# 4xx 错误不重试（key 无效等）
274+
printf "BALANCE=0\nBALANCE_NUM=0\nTOTAL=0\nUSED=0\nEXPIRES=Invalid key\nEXPIRES_FULL=Invalid key\nEXPIRES_TS=\nRAW=http_%s\n" "$http_status"
275+
return
276+
fi
277+
fi
278+
# 重试前等待
279+
(( attempt < DKM_CURL_RETRIES )) && sleep 0.5
280+
done
269281

282+
# 所有重试都失败
270283
if [ "$http_status" != "200" ]; then
271-
printf "BALANCE=0\nBALANCE_NUM=0\nTOTAL=0\nUSED=0\nEXPIRES=Invalid key\nEXPIRES_FULL=Invalid key\nEXPIRES_TS=\nRAW=http_%s\n" "$http_status"
284+
printf "BALANCE=0\nBALANCE_NUM=0\nTOTAL=0\nUSED=0\nEXPIRES=?\nEXPIRES_FULL=?\nEXPIRES_TS=\nRAW=curl_error\n"
272285
return
273286
fi
274287

‎bin/dk.ps1‎

Lines changed: 134 additions & 106 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,8 @@ $script:DK_PATH = if ($PSCommandPath) { $PSCommandPath } elseif ($MyInvocation.M
2727
$script:DK_DIR = if ($script:DK_PATH) { Split-Path $script:DK_PATH -Parent } else { $null }
2828
$script:SALT = "oroio"
2929
$script:CACHE_TTL = 30
30-
$script:CURL_TIMEOUT = 4
30+
$script:CURL_TIMEOUT = 8
31+
$script:CURL_RETRIES = 3
3132
$script:LIST_MAX_JOBS = 6
3233

3334
function Show-Usage {
@@ -388,78 +389,93 @@ function Fetch-Usage {
388389
RAW = ""
389390
}
390391

391-
try {
392-
$headers = @{
393-
"Authorization" = "Bearer $Key"
394-
"User-Agent" = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
395-
}
396-
397-
$response = Invoke-RestMethod -Uri "https://app.factory.ai/api/organization/members/chat-usage" `
398-
-Headers $headers -Method Get -TimeoutSec $script:CURL_TIMEOUT -ErrorAction Stop
399-
400-
$usage = $response.usage
401-
if ($null -eq $usage) {
402-
$result.RAW = "no_usage"
403-
return $result
404-
}
405-
406-
$section = $null
407-
foreach ($s in @($usage.standard, $usage.premium, $usage.total, $usage.main)) {
408-
if ($null -ne $s) {
409-
$section = $s
410-
break
411-
}
412-
}
413-
414-
if ($null -ne $section) {
415-
$total = $section.totalAllowance
416-
if ($null -eq $total) { $total = $section.basicAllowance }
417-
if ($null -eq $total) { $total = $section.allowance }
392+
$headers = @{
393+
"Authorization" = "Bearer $Key"
394+
"User-Agent" = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
395+
}
396+
397+
for ($attempt = 1; $attempt -le $script:CURL_RETRIES; $attempt++) {
398+
try {
399+
$response = Invoke-RestMethod -Uri "https://app.factory.ai/api/organization/members/chat-usage" `
400+
-Headers $headers -Method Get -TimeoutSec $script:CURL_TIMEOUT -ErrorAction Stop
418401

419-
$used = $section.orgTotalTokensUsed
420-
if ($null -eq $used) { $used = $section.used }
421-
if ($null -eq $used) { $used = $section.tokensUsed }
422-
if ($null -eq $used) { $used = 0 }
402+
$usage = $response.usage
403+
if ($null -eq $usage) {
404+
$result.RAW = "no_usage"
405+
return $result
406+
}
423407

424-
$overage = $section.orgOverageUsed
425-
if ($null -eq $overage) { $overage = 0 }
426-
$used = $used + $overage
408+
$section = $null
409+
foreach ($s in @($usage.standard, $usage.premium, $usage.total, $usage.main)) {
410+
if ($null -ne $s) {
411+
$section = $s
412+
break
413+
}
414+
}
427415

428-
if ($null -ne $total) {
429-
$result.TOTAL = [long]$total
430-
$result.USED = [long]$used
431-
$result.BALANCE_NUM = [long]($total - $used)
432-
$result.BALANCE = $result.BALANCE_NUM
416+
if ($null -ne $section) {
417+
$total = $section.totalAllowance
418+
if ($null -eq $total) { $total = $section.basicAllowance }
419+
if ($null -eq $total) { $total = $section.allowance }
420+
421+
$used = $section.orgTotalTokensUsed
422+
if ($null -eq $used) { $used = $section.used }
423+
if ($null -eq $used) { $used = $section.tokensUsed }
424+
if ($null -eq $used) { $used = 0 }
425+
426+
$overage = $section.orgOverageUsed
427+
if ($null -eq $overage) { $overage = 0 }
428+
$used = $used + $overage
429+
430+
if ($null -ne $total) {
431+
$result.TOTAL = [long]$total
432+
$result.USED = [long]$used
433+
$result.BALANCE_NUM = [long]($total - $used)
434+
$result.BALANCE = $result.BALANCE_NUM
435+
}
433436
}
434-
}
435-
436-
$expRaw = $usage.endDate
437-
if ($null -eq $expRaw) { $expRaw = $usage.expire_at }
438-
if ($null -eq $expRaw) { $expRaw = $usage.expires_at }
439-
440-
if ($null -ne $expRaw) {
441-
try {
442-
if ($expRaw -match '^\d+$') {
443-
$ts = [long]$expRaw / 1000
444-
$date = [DateTimeOffset]::FromUnixTimeSeconds([long]$ts)
445-
$result.EXPIRES = $date.ToString("yyyy-MM-dd")
437+
438+
$expRaw = $usage.endDate
439+
if ($null -eq $expRaw) { $expRaw = $usage.expire_at }
440+
if ($null -eq $expRaw) { $expRaw = $usage.expires_at }
441+
442+
if ($null -ne $expRaw) {
443+
try {
444+
if ($expRaw -match '^\d+$') {
445+
$ts = [long]$expRaw / 1000
446+
$date = [DateTimeOffset]::FromUnixTimeSeconds([long]$ts)
447+
$result.EXPIRES = $date.ToString("yyyy-MM-dd")
448+
}
449+
else {
450+
$result.EXPIRES = $expRaw.ToString()
451+
}
446452
}
447-
else {
453+
catch {
448454
$result.EXPIRES = $expRaw.ToString()
449455
}
450456
}
451-
catch {
452-
$result.EXPIRES = $expRaw.ToString()
457+
return $result
458+
}
459+
catch [System.Net.WebException] {
460+
$statusCode = [int]$_.Exception.Response.StatusCode
461+
if ($statusCode -ge 400 -and $statusCode -lt 500) {
462+
$result.RAW = "http_$statusCode"
463+
$result.EXPIRES = "Invalid key"
464+
return $result
465+
}
466+
if ($attempt -lt $script:CURL_RETRIES) {
467+
Start-Sleep -Milliseconds 500
468+
}
469+
}
470+
catch {
471+
if ($attempt -lt $script:CURL_RETRIES) {
472+
Start-Sleep -Milliseconds 500
453473
}
454474
}
455-
}
456-
catch {
457-
$result.BALANCE = 0
458-
$result.BALANCE_NUM = 0
459-
$result.RAW = "http_error"
460-
$result.EXPIRES = "Invalid key"
461475
}
462476

477+
$result.RAW = "http_error"
478+
$result.EXPIRES = "Invalid key"
463479
return $result
464480
}
465481

@@ -546,61 +562,73 @@ function Fetch-UsageParallel {
546562
param([string[]]$Keys)
547563

548564
$timeout = $script:CURL_TIMEOUT
565+
$retries = $script:CURL_RETRIES
549566
$maxJobs = $script:LIST_MAX_JOBS
550567
if ($maxJobs -lt 1) { $maxJobs = 6 }
551568
if ($Keys.Length -lt $maxJobs) { $maxJobs = $Keys.Length }
552569

553570
$scriptBlock = {
554-
param([string]$Key, [int]$Timeout)
571+
param([string]$Key, [int]$Timeout, [int]$Retries)
555572
$result = @{
556573
BALANCE = 0; BALANCE_NUM = 0; TOTAL = 0; USED = 0; EXPIRES = "?"; RAW = ""
557574
}
558-
try {
559-
$headers = @{
560-
"Authorization" = "Bearer $Key"
561-
"User-Agent" = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
562-
}
563-
$response = Invoke-RestMethod -Uri "https://app.factory.ai/api/organization/members/chat-usage" `
564-
-Headers $headers -Method Get -TimeoutSec $Timeout -ErrorAction Stop
565-
$usage = $response.usage
566-
if ($null -eq $usage) { $result.RAW = "no_usage"; return $result }
567-
$section = $null
568-
foreach ($s in @($usage.standard, $usage.premium, $usage.total, $usage.main)) {
569-
if ($null -ne $s) { $section = $s; break }
570-
}
571-
if ($null -ne $section) {
572-
$total = $section.totalAllowance
573-
if ($null -eq $total) { $total = $section.basicAllowance }
574-
if ($null -eq $total) { $total = $section.allowance }
575-
$used = $section.orgTotalTokensUsed
576-
if ($null -eq $used) { $used = $section.used }
577-
if ($null -eq $used) { $used = $section.tokensUsed }
578-
if ($null -eq $used) { $used = 0 }
579-
$overage = $section.orgOverageUsed
580-
if ($null -eq $overage) { $overage = 0 }
581-
$used = $used + $overage
582-
if ($null -ne $total) {
583-
$result.TOTAL = [long]$total
584-
$result.USED = [long]$used
585-
$result.BALANCE_NUM = [long]($total - $used)
586-
$result.BALANCE = $result.BALANCE_NUM
575+
$headers = @{
576+
"Authorization" = "Bearer $Key"
577+
"User-Agent" = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
578+
}
579+
for ($attempt = 1; $attempt -le $Retries; $attempt++) {
580+
try {
581+
$response = Invoke-RestMethod -Uri "https://app.factory.ai/api/organization/members/chat-usage" `
582+
-Headers $headers -Method Get -TimeoutSec $Timeout -ErrorAction Stop
583+
$usage = $response.usage
584+
if ($null -eq $usage) { $result.RAW = "no_usage"; return $result }
585+
$section = $null
586+
foreach ($s in @($usage.standard, $usage.premium, $usage.total, $usage.main)) {
587+
if ($null -ne $s) { $section = $s; break }
587588
}
589+
if ($null -ne $section) {
590+
$total = $section.totalAllowance
591+
if ($null -eq $total) { $total = $section.basicAllowance }
592+
if ($null -eq $total) { $total = $section.allowance }
593+
$used = $section.orgTotalTokensUsed
594+
if ($null -eq $used) { $used = $section.used }
595+
if ($null -eq $used) { $used = $section.tokensUsed }
596+
if ($null -eq $used) { $used = 0 }
597+
$overage = $section.orgOverageUsed
598+
if ($null -eq $overage) { $overage = 0 }
599+
$used = $used + $overage
600+
if ($null -ne $total) {
601+
$result.TOTAL = [long]$total
602+
$result.USED = [long]$used
603+
$result.BALANCE_NUM = [long]($total - $used)
604+
$result.BALANCE = $result.BALANCE_NUM
605+
}
606+
}
607+
$expRaw = $usage.endDate
608+
if ($null -eq $expRaw) { $expRaw = $usage.expire_at }
609+
if ($null -eq $expRaw) { $expRaw = $usage.expires_at }
610+
if ($null -ne $expRaw) {
611+
try {
612+
if ($expRaw -match '^\d+$') {
613+
$ts = [long]$expRaw / 1000
614+
$date = [DateTimeOffset]::FromUnixTimeSeconds([long]$ts)
615+
$result.EXPIRES = $date.ToString("yyyy-MM-dd")
616+
} else { $result.EXPIRES = $expRaw.ToString() }
617+
} catch { $result.EXPIRES = $expRaw.ToString() }
618+
}
619+
return $result
620+
} catch [System.Net.WebException] {
621+
$statusCode = [int]$_.Exception.Response.StatusCode
622+
if ($statusCode -ge 400 -and $statusCode -lt 500) {
623+
$result.RAW = "http_$statusCode"; $result.EXPIRES = "Invalid key"
624+
return $result
625+
}
626+
if ($attempt -lt $Retries) { Start-Sleep -Milliseconds 500 }
627+
} catch {
628+
if ($attempt -lt $Retries) { Start-Sleep -Milliseconds 500 }
588629
}
589-
$expRaw = $usage.endDate
590-
if ($null -eq $expRaw) { $expRaw = $usage.expire_at }
591-
if ($null -eq $expRaw) { $expRaw = $usage.expires_at }
592-
if ($null -ne $expRaw) {
593-
try {
594-
if ($expRaw -match '^\d+$') {
595-
$ts = [long]$expRaw / 1000
596-
$date = [DateTimeOffset]::FromUnixTimeSeconds([long]$ts)
597-
$result.EXPIRES = $date.ToString("yyyy-MM-dd")
598-
} else { $result.EXPIRES = $expRaw.ToString() }
599-
} catch { $result.EXPIRES = $expRaw.ToString() }
600-
}
601-
} catch {
602-
$result.BALANCE = 0; $result.BALANCE_NUM = 0; $result.RAW = "http_error"; $result.EXPIRES = "Invalid key"
603630
}
631+
$result.RAW = "http_error"; $result.EXPIRES = "Invalid key"
604632
return $result
605633
}
606634

@@ -611,7 +639,7 @@ function Fetch-UsageParallel {
611639
for ($i = 0; $i -lt $Keys.Length; $i++) {
612640
$ps = [powershell]::Create()
613641
$ps.RunspacePool = $runspacePool
614-
[void]$ps.AddScript($scriptBlock).AddArgument($Keys[$i]).AddArgument($timeout)
642+
[void]$ps.AddScript($scriptBlock).AddArgument($Keys[$i]).AddArgument($timeout).AddArgument($retries)
615643
$jobs += @{ Index = $i; PS = $ps; Handle = $ps.BeginInvoke() }
616644
}
617645

0 commit comments

Comments
 (0)