Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
2567 lines
90 KiB
Go
2567 lines
90 KiB
Go
package service
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"log/slog"
|
||
"net/http"
|
||
"regexp"
|
||
"strconv"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||
"github.com/tidwall/gjson"
|
||
)
|
||
|
||
// RateLimitService 处理限流和过载状态管理
|
||
type RateLimitService struct {
|
||
accountRepo AccountRepository
|
||
usageRepo UsageLogRepository
|
||
cfg *config.Config
|
||
geminiQuotaService *GeminiQuotaService
|
||
tempUnschedCache TempUnschedCache
|
||
timeoutCounterCache TimeoutCounterCache
|
||
openAI403CounterCache OpenAI403CounterCache
|
||
settingService *SettingService
|
||
tokenCacheInvalidator TokenCacheInvalidator
|
||
runtimeBlocker AccountRuntimeBlocker
|
||
usageCacheMu sync.RWMutex
|
||
usageCache map[int64]*geminiUsageCacheEntry
|
||
|
||
// OpenAI Team 联动熔断的进程内去重:teamID → 去重窗口截止时间
|
||
openaiTeamLinkedMu sync.Mutex
|
||
openaiTeamLinkedRecent map[string]time.Time
|
||
}
|
||
|
||
type AccountRuntimeBlocker interface {
|
||
BlockAccountScheduling(account *Account, until time.Time, reason string)
|
||
ClearAccountSchedulingBlock(accountID int64)
|
||
}
|
||
|
||
// SuccessfulTestRecoveryResult 表示测试成功后恢复了哪些运行时状态。
|
||
type SuccessfulTestRecoveryResult struct {
|
||
ClearedError bool
|
||
ClearedRateLimit bool
|
||
}
|
||
|
||
// AccountRecoveryOptions 控制账号恢复时的附加行为。
|
||
type AccountRecoveryOptions struct {
|
||
InvalidateToken bool
|
||
}
|
||
|
||
type geminiUsageCacheEntry struct {
|
||
windowStart time.Time
|
||
cachedAt time.Time
|
||
totals GeminiUsageTotals
|
||
}
|
||
|
||
type geminiUsageTotalsBatchProvider interface {
|
||
GetGeminiUsageTotalsBatch(ctx context.Context, accountIDs []int64, startTime, endTime time.Time) (map[int64]GeminiUsageTotals, error)
|
||
}
|
||
|
||
const geminiPrecheckCacheTTL = time.Minute
|
||
|
||
const (
|
||
defaultRateLimit429CooldownSeconds = 5
|
||
maxRateLimit429CooldownSeconds = 7200
|
||
)
|
||
|
||
const (
|
||
openAIImageRateLimitDefaultCooldown = time.Minute
|
||
openAIImageRateLimitReason = "openai_image_rate_limited"
|
||
)
|
||
|
||
var openAIImageTryAgainPattern = regexp.MustCompile(`(?i)try again in\s+([0-9]+(?:\.[0-9]+)?)\s*(ms|s|sec|secs|second|seconds|m|min|mins|minute|minutes)`)
|
||
|
||
const (
|
||
openAI403CooldownMinutesDefault = 10
|
||
openAI403DisableThreshold = 3
|
||
openAI403CounterWindowMinutes = 180
|
||
)
|
||
|
||
// NewRateLimitService 创建RateLimitService实例
|
||
func NewRateLimitService(accountRepo AccountRepository, usageRepo UsageLogRepository, cfg *config.Config, geminiQuotaService *GeminiQuotaService, tempUnschedCache TempUnschedCache) *RateLimitService {
|
||
return &RateLimitService{
|
||
accountRepo: accountRepo,
|
||
usageRepo: usageRepo,
|
||
cfg: cfg,
|
||
geminiQuotaService: geminiQuotaService,
|
||
tempUnschedCache: tempUnschedCache,
|
||
usageCache: make(map[int64]*geminiUsageCacheEntry),
|
||
}
|
||
}
|
||
|
||
// SetTimeoutCounterCache 设置超时计数器缓存(可选依赖)
|
||
func (s *RateLimitService) SetTimeoutCounterCache(cache TimeoutCounterCache) {
|
||
s.timeoutCounterCache = cache
|
||
}
|
||
|
||
// SetOpenAI403CounterCache 设置 OpenAI 403 连续失败计数器(可选依赖)
|
||
func (s *RateLimitService) SetOpenAI403CounterCache(cache OpenAI403CounterCache) {
|
||
s.openAI403CounterCache = cache
|
||
}
|
||
|
||
// SetSettingService 设置系统设置服务(可选依赖)
|
||
func (s *RateLimitService) SetSettingService(settingService *SettingService) {
|
||
s.settingService = settingService
|
||
}
|
||
|
||
// SetTokenCacheInvalidator 设置 token 缓存清理器(可选依赖)
|
||
func (s *RateLimitService) SetTokenCacheInvalidator(invalidator TokenCacheInvalidator) {
|
||
s.tokenCacheInvalidator = invalidator
|
||
}
|
||
|
||
func (s *RateLimitService) SetAccountRuntimeBlocker(blocker AccountRuntimeBlocker) {
|
||
s.runtimeBlocker = blocker
|
||
}
|
||
|
||
func (s *RateLimitService) IsOpenAIAdvancedSchedulerStickyWeightedEnabled(ctx context.Context) bool {
|
||
if s == nil || s.settingService == nil {
|
||
return false
|
||
}
|
||
gateway := &OpenAIGatewayService{rateLimitService: s}
|
||
return gateway.isOpenAIAdvancedSchedulerStickyWeightedEnabled(ctx)
|
||
}
|
||
|
||
func (s *RateLimitService) notifyAccountSchedulingBlocked(account *Account, until time.Time, reason string) {
|
||
if s == nil || s.runtimeBlocker == nil || account == nil {
|
||
return
|
||
}
|
||
s.runtimeBlocker.BlockAccountScheduling(account, until, reason)
|
||
}
|
||
|
||
func (s *RateLimitService) notifyAccountSchedulingBlockCleared(accountID int64) {
|
||
if s == nil || s.runtimeBlocker == nil || accountID <= 0 {
|
||
return
|
||
}
|
||
s.runtimeBlocker.ClearAccountSchedulingBlock(accountID)
|
||
}
|
||
|
||
// ApplyAccountSchedulingThreshold evaluates admin-configured per-platform
|
||
// utilization thresholds and, when breached, parks the account as temp-
|
||
// unschedulable until the winning window resets. Returns true when the account
|
||
// is blocked (either newly or already paused for the same threshold reason).
|
||
func (s *RateLimitService) ApplyAccountSchedulingThreshold(ctx context.Context, account *Account) bool {
|
||
if s == nil || s.settingService == nil || s.accountRepo == nil || account == nil || account.ID <= 0 {
|
||
return false
|
||
}
|
||
if !account.IsActive() || !account.Schedulable {
|
||
return false
|
||
}
|
||
|
||
now := time.Now().UTC()
|
||
thresholds := s.settingService.GetAccountSchedulingThresholds(ctx)
|
||
decision := EvaluateAccountSchedulingThreshold(account, thresholds, now)
|
||
if !decision.ShouldPause || decision.Until == nil || !decision.Until.After(now) {
|
||
return false
|
||
}
|
||
|
||
reason := BuildDetailedAccountSchedulingThresholdReason(AccountSchedulingThresholdReasonInput{
|
||
Platform: decision.Platform,
|
||
Window: decision.Window,
|
||
Scope: decision.Scope,
|
||
ThresholdPercent: decision.ThresholdPercent,
|
||
UsedPercent: decision.UsedPercent,
|
||
Until: *decision.Until,
|
||
Now: now,
|
||
})
|
||
|
||
if accountHasSameSchedulingThresholdPause(account, *decision.Until, reason) {
|
||
return true
|
||
}
|
||
if !account.IsSchedulable() {
|
||
return false
|
||
}
|
||
|
||
account.TempUnschedulableUntil = cloneTimePtr(decision.Until)
|
||
account.TempUnschedulableReason = reason
|
||
s.notifyAccountSchedulingBlocked(account, *decision.Until, "account_scheduling_threshold")
|
||
|
||
if err := s.accountRepo.SetTempUnschedulable(ctx, account.ID, *decision.Until, reason); err != nil {
|
||
slog.Warn("account_scheduling_threshold_set_temp_unsched_failed",
|
||
"account_id", account.ID,
|
||
"platform", decision.Platform,
|
||
"window", decision.Window,
|
||
"scope", decision.Scope,
|
||
"threshold_percent", decision.ThresholdPercent,
|
||
"used_percent", decision.UsedPercent,
|
||
"until", decision.Until.UTC(),
|
||
"error", err)
|
||
} else if s.tempUnschedCache != nil {
|
||
if state := tempUnschedStateFromStoredReason(reason, decision.Until.Unix()); state != nil {
|
||
if err := s.tempUnschedCache.SetTempUnsched(ctx, account.ID, state); err != nil {
|
||
slog.Warn("account_scheduling_threshold_cache_set_failed", "account_id", account.ID, "error", err)
|
||
}
|
||
}
|
||
}
|
||
|
||
slog.Info("account_scheduling_threshold_temp_unschedulable",
|
||
"account_id", account.ID,
|
||
"platform", decision.Platform,
|
||
"window", decision.Window,
|
||
"scope", decision.Scope,
|
||
"threshold_percent", decision.ThresholdPercent,
|
||
"used_percent", decision.UsedPercent,
|
||
"until", decision.Until.UTC())
|
||
return true
|
||
}
|
||
|
||
func accountHasSameSchedulingThresholdPause(account *Account, until time.Time, reason string) bool {
|
||
if account == nil || account.TempUnschedulableUntil == nil {
|
||
return false
|
||
}
|
||
if account.TempUnschedulableUntil.UTC().Unix() != until.UTC().Unix() {
|
||
return false
|
||
}
|
||
|
||
existing, ok := parseTempUnschedReasonPayload(account.TempUnschedulableReason)
|
||
if !ok || existing.Source != AccountSchedulingThresholdReasonSource {
|
||
return false
|
||
}
|
||
next, ok := parseTempUnschedReasonPayload(reason)
|
||
if !ok || next.Source != AccountSchedulingThresholdReasonSource {
|
||
return false
|
||
}
|
||
|
||
existing.TriggeredAtUnix = 0
|
||
next.TriggeredAtUnix = 0
|
||
return existing == next
|
||
}
|
||
|
||
// ErrorPolicyResult 表示错误策略检查的结果
|
||
type ErrorPolicyResult int
|
||
|
||
const (
|
||
ErrorPolicyNone ErrorPolicyResult = iota // 未命中任何策略,继续默认逻辑
|
||
ErrorPolicySkipped // 自定义错误码开启但未命中,跳过处理
|
||
ErrorPolicyMatched // 自定义错误码命中,应停止调度
|
||
ErrorPolicyTempUnscheduled // 临时不可调度规则命中
|
||
)
|
||
|
||
// CheckErrorPolicy 检查自定义错误码和临时不可调度规则。
|
||
// 自定义错误码开启时覆盖后续所有逻辑(包括临时不可调度)。
|
||
func (s *RateLimitService) CheckErrorPolicy(ctx context.Context, account *Account, statusCode int, responseBody []byte, requestedModel ...string) ErrorPolicyResult {
|
||
ctx = withTempUnschedulableModel(ctx, requestedModel)
|
||
if account.IsCustomErrorCodesEnabled() {
|
||
if account.ShouldHandleErrorCode(statusCode) {
|
||
return ErrorPolicyMatched
|
||
}
|
||
slog.Info("account_error_code_skipped", "account_id", account.ID, "status_code", statusCode)
|
||
return ErrorPolicySkipped
|
||
}
|
||
if account.IsPoolMode() {
|
||
// 池模式只跳过默认账号状态处理;管理员显式配置的临时不可调度规则仍应生效。
|
||
// 401 保留现有认证错误语义,避免改变重复 401 的升级行为。
|
||
if statusCode != http.StatusUnauthorized && s.tryTempUnschedulable(ctx, account, statusCode, responseBody) {
|
||
return ErrorPolicyTempUnscheduled
|
||
}
|
||
return ErrorPolicySkipped
|
||
}
|
||
if s.tryTempUnschedulable(ctx, account, statusCode, responseBody, firstRequestedModel(requestedModel)) {
|
||
return ErrorPolicyTempUnscheduled
|
||
}
|
||
return ErrorPolicyNone
|
||
}
|
||
|
||
// HandleUpstreamError 处理上游错误响应,标记账号状态
|
||
// 返回是否应该停止该账号的调度
|
||
func (s *RateLimitService) HandleUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte, requestedModel ...string) (shouldDisable bool) {
|
||
ctx = withTempUnschedulableModel(ctx, requestedModel)
|
||
// Team 联动熔断必须先于池模式/自定义错误码/临时不可调度的各类早退;
|
||
// 同请求内与 fastpath 调用点的重复触发由方法内去重吸收。
|
||
s.maybeHandleOpenAITeamLinkedError(ctx, account, statusCode, responseBody)
|
||
customErrorCodesEnabled := account.IsCustomErrorCodesEnabled()
|
||
|
||
// 池模式默认不标记本地账号状态;但管理员显式配置的临时不可调度规则优先。
|
||
// 401 保留现有认证错误语义,不在这里改变池模式的认证处理。
|
||
if account.IsPoolMode() && !customErrorCodesEnabled {
|
||
if statusCode != http.StatusUnauthorized && s.tryTempUnschedulable(ctx, account, statusCode, responseBody) {
|
||
return true
|
||
}
|
||
slog.Info("pool_mode_error_skipped", "account_id", account.ID, "status_code", statusCode)
|
||
return false
|
||
}
|
||
|
||
// apikey 类型账号:检查自定义错误码配置
|
||
// 如果启用且错误码不在列表中,则不处理(不停止调度、不标记限流/过载)
|
||
if !account.ShouldHandleErrorCode(statusCode) {
|
||
slog.Info("account_error_code_skipped", "account_id", account.ID, "status_code", statusCode)
|
||
return false
|
||
}
|
||
|
||
if len(requestedModel) > 0 && s.HandleUpstreamModelNotFound(ctx, account, requestedModel[0], statusCode, responseBody) {
|
||
return true
|
||
}
|
||
|
||
// Anthropic official 5h / 7d window exhaustion is a hard account limit.
|
||
// It must take precedence over user-configured 429 temp-unsched rules,
|
||
// otherwise a broad "rate limit" keyword rule can shorten a multi-hour
|
||
// cooldown to a local temporary pause.
|
||
if statusCode == http.StatusTooManyRequests && account.Platform == PlatformAnthropic {
|
||
// 7d_oi 是 Fable 模型专属的 7d 窗口:只标记模型级限流,账号对其他模型仍可调度。
|
||
fableLimited := s.persistAnthropicFableWindowLimit(ctx, account, headers)
|
||
if s.persistAnthropicExhaustedWindowLimit(ctx, account, headers) {
|
||
return false
|
||
}
|
||
if fableLimited {
|
||
return false
|
||
}
|
||
}
|
||
|
||
// 先尝试临时不可调度规则(401除外)
|
||
// 如果匹配成功,直接返回,不执行后续禁用逻辑
|
||
if statusCode != 401 {
|
||
if s.tryTempUnschedulable(ctx, account, statusCode, responseBody, firstRequestedModel(requestedModel)) {
|
||
return true
|
||
}
|
||
}
|
||
|
||
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(responseBody))
|
||
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
|
||
if upstreamMsg != "" {
|
||
upstreamMsg = truncateForLog([]byte(upstreamMsg), 512)
|
||
}
|
||
|
||
switch statusCode {
|
||
case 400:
|
||
// "organization has been disabled" → 永久禁用
|
||
if strings.Contains(strings.ToLower(upstreamMsg), "organization has been disabled") {
|
||
msg := "Organization disabled (400): " + upstreamMsg
|
||
s.handleAuthError(ctx, account, msg)
|
||
shouldDisable = true
|
||
} else if account.Platform == PlatformAnthropic && strings.Contains(strings.ToLower(upstreamMsg), "credit balance") {
|
||
// Anthropic API key 余额不足(语义等同 402),停止调度
|
||
msg := "Credit balance exhausted (400): " + upstreamMsg
|
||
s.handleAuthError(ctx, account, msg)
|
||
shouldDisable = true
|
||
} else if strings.Contains(strings.ToLower(upstreamMsg), "identity verification is required") {
|
||
// KYC 身份验证要求 → 永久禁用,账号需完成身份验证后才能恢复
|
||
msg := "Identity verification required (400): " + upstreamMsg
|
||
s.handleAuthError(ctx, account, msg)
|
||
shouldDisable = true
|
||
}
|
||
// 其他 400 错误(如参数问题)不处理,不禁用账号
|
||
case 401:
|
||
// 外审第9轮:Spark 影子无独立凭据,401 是母账号 token 问题——失效缓存 / refresh_token 判断 /
|
||
// 永久禁用 / 临时不可调度都必须落到凭据 owner(母账号),否则影子(无 refresh_token)必中
|
||
// "refresh_token missing"永久禁用分支、母账号 token cache 也不会被清,把母账号可恢复的 token
|
||
// 问题变成影子永久死亡。母账号被标记 temp-unschedulable 后由 parentHealthyForShadow 级联排除影子。
|
||
// 非影子时 resolveCredentialAccount 返回自身;母账号缺失/损坏(orphan 影子,罕见)时回退到原 account。
|
||
authAccount := account
|
||
if resolved, rerr := resolveCredentialAccount(ctx, s.accountRepo, account); rerr == nil && resolved != nil {
|
||
authAccount = resolved
|
||
}
|
||
// OpenAI: token_invalidated / token_revoked 表示 token 被永久作废(非过期),直接标记 error
|
||
openai401Code := extractUpstreamErrorCode(responseBody)
|
||
if authAccount.Platform == PlatformOpenAI && (openai401Code == "token_invalidated" || openai401Code == "token_revoked") {
|
||
msg := "Token revoked (401): account authentication permanently revoked"
|
||
if upstreamMsg != "" {
|
||
msg = "Token revoked (401): " + upstreamMsg
|
||
}
|
||
s.handleAuthError(ctx, authAccount, msg)
|
||
shouldDisable = true
|
||
break
|
||
}
|
||
// OpenAI: {"detail":"Unauthorized"} 表示 token 完全无效(非标准 OpenAI 错误格式),直接标记 error
|
||
if authAccount.Platform == PlatformOpenAI && gjson.GetBytes(responseBody, "detail").String() == "Unauthorized" {
|
||
msg := "Unauthorized (401): account authentication failed permanently"
|
||
if upstreamMsg != "" {
|
||
msg = "Unauthorized (401): " + upstreamMsg
|
||
}
|
||
s.handleAuthError(ctx, authAccount, msg)
|
||
shouldDisable = true
|
||
break
|
||
}
|
||
// OAuth 账号在 401 错误时临时不可调度(给 token 刷新窗口);非 OAuth 账号保持原有 SetError 行为。
|
||
if authAccount.Type == AccountTypeOAuth {
|
||
// 1. 失效缓存
|
||
if s.tokenCacheInvalidator != nil {
|
||
if err := s.tokenCacheInvalidator.InvalidateToken(ctx, authAccount); err != nil {
|
||
slog.Warn("oauth_401_invalidate_cache_failed", "account_id", authAccount.ID, "error", err)
|
||
}
|
||
}
|
||
// 缺少 refresh_token 的 OAuth 账号无法在冷却期内自愈(后台刷新服务也会跳过),
|
||
// 直接走 SetError 永久禁用,避免冷却结束后再被选中产生一发无意义的 502。
|
||
if strings.TrimSpace(authAccount.GetCredential("refresh_token")) == "" {
|
||
msg := "Authentication failed (401): refresh_token missing, cannot recover"
|
||
if upstreamMsg != "" {
|
||
msg = "OAuth 401 (no refresh_token): " + upstreamMsg
|
||
}
|
||
s.handleAuthError(ctx, authAccount, msg)
|
||
shouldDisable = true
|
||
break
|
||
}
|
||
// 2. 临时不可调度,替代 SetError(保持 status=active 让刷新服务能拾取)
|
||
// 注意:此处不再写回 account.Credentials/expires_at。
|
||
// 原实现使用请求开始时的 account 快照整列覆盖 credentials JSONB(见
|
||
// persistAccountCredentials → accountRepository.UpdateCredentials → SetCredentials),
|
||
// 在另一个 worker 刚刷新完 refresh_token 的窄窗口内会把新 refresh_token 回滚为旧值,
|
||
// 导致下一周期用旧 refresh_token 调上游拿到 invalid_grant 后,
|
||
// tryRecoverFromRefreshRace 重读 DB 发现 currentRT == usedRT 也救不回来,账号被错误 disable。
|
||
// 这里仅依赖 InvalidateToken + SetTempUnschedulable 让账号在冷却期内不被调度,
|
||
// 冷却结束后由 token_provider 的 NeedsRefresh / token_refresh_service 走带分布式锁的正路刷新。
|
||
msg := "Authentication failed (401): invalid or expired credentials"
|
||
if upstreamMsg != "" {
|
||
msg = "OAuth 401: " + upstreamMsg
|
||
}
|
||
if authAccount.Platform == PlatformAntigravity {
|
||
extraUpdates := antigravityForceTokenRefreshExtra("401_invalid")
|
||
if err := s.accountRepo.UpdateExtra(ctx, authAccount.ID, extraUpdates); err != nil {
|
||
slog.Warn("antigravity_401_force_refresh_mark_failed", "account_id", authAccount.ID, "error", err)
|
||
} else {
|
||
if authAccount.Extra == nil {
|
||
authAccount.Extra = make(map[string]any, len(extraUpdates))
|
||
}
|
||
for k, v := range extraUpdates {
|
||
authAccount.Extra[k] = v
|
||
}
|
||
slog.Info("antigravity_401_force_refresh_marked", "account_id", authAccount.ID)
|
||
}
|
||
}
|
||
cooldownMinutes := s.cfg.RateLimit.OAuth401CooldownMinutes
|
||
if cooldownMinutes <= 0 {
|
||
cooldownMinutes = 10
|
||
}
|
||
until := time.Now().Add(time.Duration(cooldownMinutes) * time.Minute)
|
||
s.notifyAccountSchedulingBlocked(authAccount, until, "oauth_401")
|
||
if err := s.accountRepo.SetTempUnschedulable(ctx, authAccount.ID, until, msg); err != nil {
|
||
slog.Warn("oauth_401_set_temp_unschedulable_failed", "account_id", authAccount.ID, "error", err)
|
||
}
|
||
shouldDisable = true
|
||
} else {
|
||
// 非 OAuth:保持 SetError 行为
|
||
msg := "Authentication failed (401): invalid or expired credentials"
|
||
if upstreamMsg != "" {
|
||
msg = "Authentication failed (401): " + upstreamMsg
|
||
}
|
||
s.handleAuthError(ctx, authAccount, msg)
|
||
shouldDisable = true
|
||
}
|
||
case 402:
|
||
// 国产供应商:余额不足是可恢复状态(充值/检测恢复后由周期任务自动解除),
|
||
// 不能走 handleAuthError 永久置 status=error。改为可恢复的临时停调。
|
||
if account.IsCNProvider() {
|
||
s.handleCNProviderInsufficientBalance(ctx, account, upstreamMsg)
|
||
shouldDisable = true
|
||
break
|
||
}
|
||
// OpenAI: deactivated_workspace 表示工作区已停用,直接标记 error
|
||
if account.Platform == PlatformOpenAI && gjson.GetBytes(responseBody, "detail.code").String() == "deactivated_workspace" {
|
||
msg := "Workspace deactivated (402): workspace has been deactivated"
|
||
s.handleAuthError(ctx, account, msg)
|
||
shouldDisable = true
|
||
break
|
||
}
|
||
// 支付要求:余额不足或计费问题,停止调度
|
||
msg := "Payment required (402): insufficient balance or billing issue"
|
||
if upstreamMsg != "" {
|
||
msg = "Payment required (402): " + upstreamMsg
|
||
}
|
||
s.handleAuthError(ctx, account, msg)
|
||
shouldDisable = true
|
||
case 403:
|
||
logger.LegacyPrintf(
|
||
"service.ratelimit",
|
||
"[HandleUpstreamErrorRaw] account_id=%d platform=%s type=%s status=403 request_id=%s cf_ray=%s upstream_msg=%s raw_body=%s",
|
||
account.ID,
|
||
account.Platform,
|
||
account.Type,
|
||
strings.TrimSpace(headers.Get("x-request-id")),
|
||
strings.TrimSpace(headers.Get("cf-ray")),
|
||
upstreamMsg,
|
||
truncateForLog(responseBody, 1024),
|
||
)
|
||
shouldDisable = s.handle403(ctx, account, upstreamMsg, responseBody)
|
||
case 429:
|
||
s.handle429(ctx, account, headers, responseBody)
|
||
shouldDisable = false
|
||
case 529:
|
||
s.handle529(ctx, account)
|
||
shouldDisable = false
|
||
default:
|
||
// 自定义错误码启用时:在列表中的错误码都应该停止调度
|
||
if customErrorCodesEnabled {
|
||
msg := "Custom error code triggered"
|
||
if upstreamMsg != "" {
|
||
msg = upstreamMsg
|
||
}
|
||
s.handleCustomErrorCode(ctx, account, statusCode, msg)
|
||
shouldDisable = true
|
||
} else if statusCode >= 500 {
|
||
// 未启用自定义错误码时:仅记录5xx错误
|
||
slog.Warn("account_upstream_error", "account_id", account.ID, "status_code", statusCode)
|
||
shouldDisable = false
|
||
}
|
||
}
|
||
|
||
return shouldDisable
|
||
}
|
||
|
||
// PreCheckUsage proactively checks local quota before dispatching a request.
|
||
// Returns false when the account should be skipped.
|
||
func (s *RateLimitService) PreCheckUsage(ctx context.Context, account *Account, requestedModel string) (bool, error) {
|
||
if account == nil || account.Platform != PlatformGemini {
|
||
return true, nil
|
||
}
|
||
if s.usageRepo == nil || s.geminiQuotaService == nil {
|
||
return true, nil
|
||
}
|
||
|
||
quota, ok := s.geminiQuotaService.QuotaForAccount(ctx, account)
|
||
if !ok {
|
||
return true, nil
|
||
}
|
||
|
||
now := time.Now()
|
||
modelClass := geminiModelClassFromName(requestedModel)
|
||
|
||
// 1) Daily quota precheck (RPD; resets at PST midnight)
|
||
{
|
||
var limit int64
|
||
if quota.SharedRPD > 0 {
|
||
limit = quota.SharedRPD
|
||
} else {
|
||
switch modelClass {
|
||
case geminiModelFlash:
|
||
limit = quota.FlashRPD
|
||
default:
|
||
limit = quota.ProRPD
|
||
}
|
||
}
|
||
|
||
if limit > 0 {
|
||
start := geminiDailyWindowStart(now)
|
||
totals, ok := s.getGeminiUsageTotals(account.ID, start, now)
|
||
if !ok {
|
||
stats, err := s.usageRepo.GetModelStatsWithFilters(ctx, start, now, 0, 0, account.ID, 0, nil, nil, nil)
|
||
if err != nil {
|
||
return true, err
|
||
}
|
||
totals = geminiAggregateUsage(stats)
|
||
s.setGeminiUsageTotals(account.ID, start, now, totals)
|
||
}
|
||
|
||
var used int64
|
||
if quota.SharedRPD > 0 {
|
||
used = totals.ProRequests + totals.FlashRequests
|
||
} else {
|
||
switch modelClass {
|
||
case geminiModelFlash:
|
||
used = totals.FlashRequests
|
||
default:
|
||
used = totals.ProRequests
|
||
}
|
||
}
|
||
|
||
if used >= limit {
|
||
resetAt := geminiDailyResetTime(now)
|
||
// NOTE:
|
||
// - This is a local precheck to reduce upstream 429s.
|
||
// - Do NOT mark the account as rate-limited here; rate_limit_reset_at should reflect real upstream 429s.
|
||
slog.Info("gemini_precheck_daily_quota_reached", "account_id", account.ID, "used", used, "limit", limit, "reset_at", resetAt)
|
||
return false, nil
|
||
}
|
||
}
|
||
}
|
||
|
||
// 2) Minute quota precheck (RPM; fixed window current minute)
|
||
{
|
||
var limit int64
|
||
if quota.SharedRPM > 0 {
|
||
limit = quota.SharedRPM
|
||
} else {
|
||
switch modelClass {
|
||
case geminiModelFlash:
|
||
limit = quota.FlashRPM
|
||
default:
|
||
limit = quota.ProRPM
|
||
}
|
||
}
|
||
|
||
if limit > 0 {
|
||
start := now.Truncate(time.Minute)
|
||
stats, err := s.usageRepo.GetModelStatsWithFilters(ctx, start, now, 0, 0, account.ID, 0, nil, nil, nil)
|
||
if err != nil {
|
||
return true, err
|
||
}
|
||
totals := geminiAggregateUsage(stats)
|
||
|
||
var used int64
|
||
if quota.SharedRPM > 0 {
|
||
used = totals.ProRequests + totals.FlashRequests
|
||
} else {
|
||
switch modelClass {
|
||
case geminiModelFlash:
|
||
used = totals.FlashRequests
|
||
default:
|
||
used = totals.ProRequests
|
||
}
|
||
}
|
||
|
||
if used >= limit {
|
||
resetAt := start.Add(time.Minute)
|
||
// Do not persist "rate limited" status from local precheck. See note above.
|
||
slog.Info("gemini_precheck_minute_quota_reached", "account_id", account.ID, "used", used, "limit", limit, "reset_at", resetAt)
|
||
return false, nil
|
||
}
|
||
}
|
||
}
|
||
|
||
return true, nil
|
||
}
|
||
|
||
// PreCheckUsageBatch performs quota precheck for multiple accounts in one request.
|
||
// Returned map value=false means the account should be skipped.
|
||
func (s *RateLimitService) PreCheckUsageBatch(ctx context.Context, accounts []*Account, requestedModel string) (map[int64]bool, error) {
|
||
result := make(map[int64]bool, len(accounts))
|
||
for _, account := range accounts {
|
||
if account == nil {
|
||
continue
|
||
}
|
||
result[account.ID] = true
|
||
}
|
||
|
||
if len(accounts) == 0 || requestedModel == "" {
|
||
return result, nil
|
||
}
|
||
if s.usageRepo == nil || s.geminiQuotaService == nil {
|
||
return result, nil
|
||
}
|
||
|
||
modelClass := geminiModelClassFromName(requestedModel)
|
||
now := time.Now()
|
||
dailyStart := geminiDailyWindowStart(now)
|
||
minuteStart := now.Truncate(time.Minute)
|
||
|
||
type quotaAccount struct {
|
||
account *Account
|
||
quota GeminiQuota
|
||
}
|
||
quotaAccounts := make([]quotaAccount, 0, len(accounts))
|
||
for _, account := range accounts {
|
||
if account == nil || account.Platform != PlatformGemini {
|
||
continue
|
||
}
|
||
quota, ok := s.geminiQuotaService.QuotaForAccount(ctx, account)
|
||
if !ok {
|
||
continue
|
||
}
|
||
quotaAccounts = append(quotaAccounts, quotaAccount{
|
||
account: account,
|
||
quota: quota,
|
||
})
|
||
}
|
||
if len(quotaAccounts) == 0 {
|
||
return result, nil
|
||
}
|
||
|
||
// 1) Daily precheck (cached + batch DB fallback)
|
||
dailyTotalsByID := make(map[int64]GeminiUsageTotals, len(quotaAccounts))
|
||
dailyMissIDs := make([]int64, 0, len(quotaAccounts))
|
||
for _, item := range quotaAccounts {
|
||
limit := geminiDailyLimit(item.quota, modelClass)
|
||
if limit <= 0 {
|
||
continue
|
||
}
|
||
accountID := item.account.ID
|
||
if totals, ok := s.getGeminiUsageTotals(accountID, dailyStart, now); ok {
|
||
dailyTotalsByID[accountID] = totals
|
||
continue
|
||
}
|
||
dailyMissIDs = append(dailyMissIDs, accountID)
|
||
}
|
||
if len(dailyMissIDs) > 0 {
|
||
totalsBatch, err := s.getGeminiUsageTotalsBatch(ctx, dailyMissIDs, dailyStart, now)
|
||
if err != nil {
|
||
return result, err
|
||
}
|
||
for _, accountID := range dailyMissIDs {
|
||
totals := totalsBatch[accountID]
|
||
dailyTotalsByID[accountID] = totals
|
||
s.setGeminiUsageTotals(accountID, dailyStart, now, totals)
|
||
}
|
||
}
|
||
for _, item := range quotaAccounts {
|
||
limit := geminiDailyLimit(item.quota, modelClass)
|
||
if limit <= 0 {
|
||
continue
|
||
}
|
||
accountID := item.account.ID
|
||
used := geminiUsedRequests(item.quota, modelClass, dailyTotalsByID[accountID], true)
|
||
if used >= limit {
|
||
resetAt := geminiDailyResetTime(now)
|
||
slog.Info("gemini_precheck_daily_quota_reached_batch", "account_id", accountID, "used", used, "limit", limit, "reset_at", resetAt)
|
||
result[accountID] = false
|
||
}
|
||
}
|
||
|
||
// 2) Minute precheck (batch DB)
|
||
minuteIDs := make([]int64, 0, len(quotaAccounts))
|
||
for _, item := range quotaAccounts {
|
||
accountID := item.account.ID
|
||
if !result[accountID] {
|
||
continue
|
||
}
|
||
if geminiMinuteLimit(item.quota, modelClass) <= 0 {
|
||
continue
|
||
}
|
||
minuteIDs = append(minuteIDs, accountID)
|
||
}
|
||
if len(minuteIDs) == 0 {
|
||
return result, nil
|
||
}
|
||
|
||
minuteTotalsByID, err := s.getGeminiUsageTotalsBatch(ctx, minuteIDs, minuteStart, now)
|
||
if err != nil {
|
||
return result, err
|
||
}
|
||
for _, item := range quotaAccounts {
|
||
accountID := item.account.ID
|
||
if !result[accountID] {
|
||
continue
|
||
}
|
||
|
||
limit := geminiMinuteLimit(item.quota, modelClass)
|
||
if limit <= 0 {
|
||
continue
|
||
}
|
||
|
||
used := geminiUsedRequests(item.quota, modelClass, minuteTotalsByID[accountID], false)
|
||
if used >= limit {
|
||
resetAt := minuteStart.Add(time.Minute)
|
||
slog.Info("gemini_precheck_minute_quota_reached_batch", "account_id", accountID, "used", used, "limit", limit, "reset_at", resetAt)
|
||
result[accountID] = false
|
||
}
|
||
}
|
||
|
||
return result, nil
|
||
}
|
||
|
||
func (s *RateLimitService) getGeminiUsageTotalsBatch(ctx context.Context, accountIDs []int64, start, end time.Time) (map[int64]GeminiUsageTotals, error) {
|
||
result := make(map[int64]GeminiUsageTotals, len(accountIDs))
|
||
if len(accountIDs) == 0 {
|
||
return result, nil
|
||
}
|
||
|
||
ids := make([]int64, 0, len(accountIDs))
|
||
seen := make(map[int64]struct{}, len(accountIDs))
|
||
for _, accountID := range accountIDs {
|
||
if accountID <= 0 {
|
||
continue
|
||
}
|
||
if _, ok := seen[accountID]; ok {
|
||
continue
|
||
}
|
||
seen[accountID] = struct{}{}
|
||
ids = append(ids, accountID)
|
||
}
|
||
if len(ids) == 0 {
|
||
return result, nil
|
||
}
|
||
|
||
if batchReader, ok := s.usageRepo.(geminiUsageTotalsBatchProvider); ok {
|
||
stats, err := batchReader.GetGeminiUsageTotalsBatch(ctx, ids, start, end)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
for _, accountID := range ids {
|
||
result[accountID] = stats[accountID]
|
||
}
|
||
return result, nil
|
||
}
|
||
|
||
for _, accountID := range ids {
|
||
stats, err := s.usageRepo.GetModelStatsWithFilters(ctx, start, end, 0, 0, accountID, 0, nil, nil, nil)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
result[accountID] = geminiAggregateUsage(stats)
|
||
}
|
||
return result, nil
|
||
}
|
||
|
||
func geminiDailyLimit(quota GeminiQuota, modelClass geminiModelClass) int64 {
|
||
if quota.SharedRPD > 0 {
|
||
return quota.SharedRPD
|
||
}
|
||
switch modelClass {
|
||
case geminiModelFlash:
|
||
return quota.FlashRPD
|
||
default:
|
||
return quota.ProRPD
|
||
}
|
||
}
|
||
|
||
func geminiMinuteLimit(quota GeminiQuota, modelClass geminiModelClass) int64 {
|
||
if quota.SharedRPM > 0 {
|
||
return quota.SharedRPM
|
||
}
|
||
switch modelClass {
|
||
case geminiModelFlash:
|
||
return quota.FlashRPM
|
||
default:
|
||
return quota.ProRPM
|
||
}
|
||
}
|
||
|
||
func geminiUsedRequests(quota GeminiQuota, modelClass geminiModelClass, totals GeminiUsageTotals, daily bool) int64 {
|
||
if daily {
|
||
if quota.SharedRPD > 0 {
|
||
return totals.ProRequests + totals.FlashRequests
|
||
}
|
||
} else {
|
||
if quota.SharedRPM > 0 {
|
||
return totals.ProRequests + totals.FlashRequests
|
||
}
|
||
}
|
||
switch modelClass {
|
||
case geminiModelFlash:
|
||
return totals.FlashRequests
|
||
default:
|
||
return totals.ProRequests
|
||
}
|
||
}
|
||
|
||
func (s *RateLimitService) getGeminiUsageTotals(accountID int64, windowStart, now time.Time) (GeminiUsageTotals, bool) {
|
||
s.usageCacheMu.RLock()
|
||
defer s.usageCacheMu.RUnlock()
|
||
|
||
if s.usageCache == nil {
|
||
return GeminiUsageTotals{}, false
|
||
}
|
||
|
||
entry, ok := s.usageCache[accountID]
|
||
if !ok || entry == nil {
|
||
return GeminiUsageTotals{}, false
|
||
}
|
||
if !entry.windowStart.Equal(windowStart) {
|
||
return GeminiUsageTotals{}, false
|
||
}
|
||
if now.Sub(entry.cachedAt) >= geminiPrecheckCacheTTL {
|
||
return GeminiUsageTotals{}, false
|
||
}
|
||
return entry.totals, true
|
||
}
|
||
|
||
func (s *RateLimitService) setGeminiUsageTotals(accountID int64, windowStart, now time.Time, totals GeminiUsageTotals) {
|
||
s.usageCacheMu.Lock()
|
||
defer s.usageCacheMu.Unlock()
|
||
if s.usageCache == nil {
|
||
s.usageCache = make(map[int64]*geminiUsageCacheEntry)
|
||
}
|
||
s.usageCache[accountID] = &geminiUsageCacheEntry{
|
||
windowStart: windowStart,
|
||
cachedAt: now,
|
||
totals: totals,
|
||
}
|
||
}
|
||
|
||
// GeminiCooldown returns the fallback cooldown duration for Gemini 429s based on tier.
|
||
func (s *RateLimitService) GeminiCooldown(ctx context.Context, account *Account) time.Duration {
|
||
if account == nil {
|
||
return 5 * time.Minute
|
||
}
|
||
if s.geminiQuotaService == nil {
|
||
return 5 * time.Minute
|
||
}
|
||
return s.geminiQuotaService.CooldownForAccount(ctx, account)
|
||
}
|
||
|
||
// handleAuthError 处理认证类错误(401/403),停止账号调度
|
||
func (s *RateLimitService) handleAuthError(ctx context.Context, account *Account, errorMsg string) {
|
||
s.notifyAccountSchedulingBlocked(account, time.Time{}, "auth_error")
|
||
if err := s.accountRepo.SetError(ctx, account.ID, errorMsg); err != nil {
|
||
slog.Warn("account_set_error_failed", "account_id", account.ID, "error", err)
|
||
return
|
||
}
|
||
slog.Warn("account_disabled_auth_error", "account_id", account.ID, "error", errorMsg)
|
||
}
|
||
|
||
func buildForbiddenErrorMessage(prefix string, upstreamMsg string, responseBody []byte, fallback string) string {
|
||
prefix = strings.TrimSpace(prefix)
|
||
if prefix != "" && !strings.HasSuffix(prefix, " ") {
|
||
prefix += " "
|
||
}
|
||
|
||
if msg := strings.TrimSpace(upstreamMsg); msg != "" {
|
||
return prefix + msg
|
||
}
|
||
|
||
rawBody := bytes.TrimSpace(responseBody)
|
||
if len(rawBody) > 0 {
|
||
if json.Valid(rawBody) {
|
||
var compact bytes.Buffer
|
||
if err := json.Compact(&compact, rawBody); err == nil {
|
||
return prefix + truncateForLog(compact.Bytes(), 512)
|
||
}
|
||
}
|
||
return prefix + truncateForLog(rawBody, 512)
|
||
}
|
||
|
||
return prefix + fallback
|
||
}
|
||
|
||
// handle403 处理 403 Forbidden 错误
|
||
// Antigravity 平台区分 validation/violation/generic 三种类型,均 SetError 永久禁用;
|
||
// 其他平台保持原有 SetError 行为。
|
||
func (s *RateLimitService) handle403(ctx context.Context, account *Account, upstreamMsg string, responseBody []byte) (shouldDisable bool) {
|
||
if account.Platform == PlatformAntigravity {
|
||
return s.handleAntigravity403(ctx, account, upstreamMsg, responseBody)
|
||
}
|
||
// 国产供应商与 openai 同口径:HTML 403(CDN/代理拦截页)不构成账号失效证据,
|
||
// 且 403 在 failover 状态集里会被逐账号重放——直接 SetError 会让一个坏请求/
|
||
// 一层坏代理连环永久禁用整组账号。走 HTML 豁免 + N 次累计 + 临时冷却。
|
||
if account.Platform == PlatformOpenAI || IsCNProvider(account.Platform) {
|
||
return s.handleOpenAI403(ctx, account, upstreamMsg, responseBody)
|
||
}
|
||
// 非 Antigravity 平台:保持原有行为
|
||
msg := buildForbiddenErrorMessage(
|
||
"Access forbidden (403):",
|
||
upstreamMsg,
|
||
responseBody,
|
||
"account may be suspended or lack permissions",
|
||
)
|
||
s.handleAuthError(ctx, account, msg)
|
||
return true
|
||
}
|
||
|
||
func (s *RateLimitService) handleOpenAI403(ctx context.Context, account *Account, upstreamMsg string, responseBody []byte) (shouldDisable bool) {
|
||
// 上游代理 / CDN 在请求到达 OpenAI API 之前就拦下时,回的是 HTML 403 页面而不是
|
||
// {"error":{...}} 结构化错误。这类响应描述的是「这条链路 / 这个端点被挡了」,
|
||
// 不构成账号凭据或权限失效的证据——例如无效的 /v1/responses 子路径(#5334)。
|
||
//
|
||
// 据此写账号状态会把请求级错误放大成账号级处罚:首次即 temp-unschedulable,
|
||
// 连续 openAI403DisableThreshold 次直接永久禁用;而 403 又在 failover 状态集里,
|
||
// 同一个坏请求会被逐个账号重放,足以把整组账号打下线。
|
||
//
|
||
// 与既有口径一致:count_tokens 路径的 isOpenAIOAuthInputTokensUnsupported 已把
|
||
// 「HTML 403 page without a structured error」按端点级响应处理;
|
||
// shouldApplyOpenAIAlphaSearchAccountErrorSideEffects 的不变式也是端点级错误
|
||
// 只换号、不写账号错误状态。这里只跳过账号处罚,不改变 failover 行为——
|
||
// 换个走不同代理的账号仍有可能成功。
|
||
if isHTMLResponse(responseBody) {
|
||
slog.Warn(
|
||
"openai_403_html_body_skips_account_penalty",
|
||
"account_id", account.ID,
|
||
"upstream_message", upstreamMsg,
|
||
)
|
||
return false
|
||
}
|
||
|
||
msg := buildForbiddenErrorMessage(
|
||
"Access forbidden (403):",
|
||
upstreamMsg,
|
||
responseBody,
|
||
"account may be suspended or lack permissions",
|
||
)
|
||
|
||
if s.openAI403CounterCache == nil {
|
||
s.handleAuthError(ctx, account, msg)
|
||
return true
|
||
}
|
||
|
||
count, err := s.openAI403CounterCache.IncrementOpenAI403Count(ctx, account.ID, openAI403CounterWindowMinutes)
|
||
if err != nil {
|
||
slog.Warn("openai_403_increment_failed", "account_id", account.ID, "error", err)
|
||
s.handleAuthError(ctx, account, msg)
|
||
return true
|
||
}
|
||
|
||
if count >= openAI403DisableThreshold {
|
||
msg = fmt.Sprintf("%s | consecutive_403=%d/%d", msg, count, openAI403DisableThreshold)
|
||
s.handleAuthError(ctx, account, msg)
|
||
return true
|
||
}
|
||
|
||
until := time.Now().Add(time.Duration(openAI403CooldownMinutesDefault) * time.Minute)
|
||
reason := fmt.Sprintf("OpenAI 403 temporary cooldown (%d/%d): %s", count, openAI403DisableThreshold, msg)
|
||
s.notifyAccountSchedulingBlocked(account, until, "openai_403_temp")
|
||
if err := s.accountRepo.SetTempUnschedulable(ctx, account.ID, until, reason); err != nil {
|
||
slog.Warn("openai_403_set_temp_unschedulable_failed", "account_id", account.ID, "error", err)
|
||
s.handleAuthError(ctx, account, msg)
|
||
return true
|
||
}
|
||
|
||
slog.Warn(
|
||
"openai_403_temp_unschedulable",
|
||
"account_id", account.ID,
|
||
"until", until,
|
||
"count", count,
|
||
"threshold", openAI403DisableThreshold,
|
||
)
|
||
return true
|
||
}
|
||
|
||
// handleAntigravity403 处理 Antigravity 平台的 403 错误
|
||
// validation(需要验证)→ 永久 SetError(需人工去 Google 验证后恢复)
|
||
// violation(违规封号)→ 永久 SetError(需人工处理)
|
||
// generic(通用禁止)→ 永久 SetError
|
||
func (s *RateLimitService) handleAntigravity403(ctx context.Context, account *Account, upstreamMsg string, responseBody []byte) (shouldDisable bool) {
|
||
fbType := classifyForbiddenType(string(responseBody))
|
||
|
||
switch fbType {
|
||
case forbiddenTypeValidation:
|
||
// VALIDATION_REQUIRED: 永久禁用,需人工去 Google 验证后手动恢复
|
||
msg := buildForbiddenErrorMessage(
|
||
"Validation required (403):",
|
||
upstreamMsg,
|
||
responseBody,
|
||
"account needs Google verification",
|
||
)
|
||
if validationURL := extractValidationURL(string(responseBody)); validationURL != "" {
|
||
msg += " | validation_url: " + validationURL
|
||
}
|
||
s.handleAuthError(ctx, account, msg)
|
||
return true
|
||
|
||
case forbiddenTypeViolation:
|
||
// 违规封号: 永久禁用,需人工处理
|
||
msg := buildForbiddenErrorMessage(
|
||
"Account violation (403):",
|
||
upstreamMsg,
|
||
responseBody,
|
||
"terms of service violation",
|
||
)
|
||
s.handleAuthError(ctx, account, msg)
|
||
return true
|
||
|
||
default:
|
||
// 通用 403: 保持原有行为
|
||
msg := buildForbiddenErrorMessage(
|
||
"Access forbidden (403):",
|
||
upstreamMsg,
|
||
responseBody,
|
||
"account may be suspended or lack permissions",
|
||
)
|
||
s.handleAuthError(ctx, account, msg)
|
||
return true
|
||
}
|
||
}
|
||
|
||
// handleCustomErrorCode 处理自定义错误码,停止账号调度
|
||
func (s *RateLimitService) handleCustomErrorCode(ctx context.Context, account *Account, statusCode int, errorMsg string) {
|
||
msg := "Custom error code " + strconv.Itoa(statusCode) + ": " + errorMsg
|
||
s.notifyAccountSchedulingBlocked(account, time.Time{}, "custom_error_code")
|
||
if err := s.accountRepo.SetError(ctx, account.ID, msg); err != nil {
|
||
slog.Warn("account_set_error_failed", "account_id", account.ID, "status_code", statusCode, "error", err)
|
||
return
|
||
}
|
||
slog.Warn("account_disabled_custom_error", "account_id", account.ID, "status_code", statusCode, "error", errorMsg)
|
||
}
|
||
|
||
// handle429 处理429限流错误
|
||
// 解析响应头获取重置时间,标记账号为限流状态
|
||
func (s *RateLimitService) handle429(ctx context.Context, account *Account, headers http.Header, responseBody []byte) {
|
||
// Spark 影子:限流/熔断状态 100% 由 QueryUsage(/wham/usage body 的 codex_bengalfox)驱动。
|
||
// /responses 的 429 携带的 x-codex-*/usage_limit_reached 是 global codex 道(plan/spec §8),
|
||
// 套到影子会把 spark 误耦合到 global 窗口——即便 spark 仍有配额也会被冷却到 global reset,
|
||
// 单影子场景直接变成无可用账号(外审第8轮 P1)。整段跳过;影子的 codex_* 仅由 account_usage 的
|
||
// QueryUsage→persistOpenAICodexProbeSnapshot 维护,枯竭由调度守卫处理。
|
||
if account.IsShadow() {
|
||
return
|
||
}
|
||
// 国产供应商(kimi/zhipu/deepseek)的 429 走专用可恢复路径:余额不足 → 临时停调,
|
||
// Coding Plan 窗口耗尽 → 冷却到快照重置点。未命中则继续默认 429 逻辑。
|
||
if account.IsCNProvider() {
|
||
if s.applyCNProviderReactive429(ctx, account, headers, responseBody) {
|
||
return
|
||
}
|
||
}
|
||
// 1. OpenAI 平台:优先尝试解析 x-codex-* 响应头(用于 rate_limit_exceeded)
|
||
if account.Platform == PlatformOpenAI {
|
||
persistOpenAI429PlanType(ctx, s.accountRepo, account, responseBody)
|
||
s.persistOpenAICodexSnapshot(ctx, account, headers)
|
||
if resetAt := s.calculateOpenAI429ResetTime(headers); resetAt != nil {
|
||
s.notifyAccountSchedulingBlocked(account, *resetAt, "429")
|
||
if err := s.accountRepo.SetRateLimited(ctx, account.ID, *resetAt); err != nil {
|
||
slog.Warn("rate_limit_set_failed", "account_id", account.ID, "error", err)
|
||
return
|
||
}
|
||
slog.Info("openai_account_rate_limited", "account_id", account.ID, "reset_at", *resetAt)
|
||
return
|
||
}
|
||
}
|
||
|
||
// 2. Anthropic 平台:尝试解析 per-window 头(5h / 7d),选择实际触发的窗口
|
||
if result := calculateAnthropic429ResetTime(headers); result != nil {
|
||
s.notifyAccountSchedulingBlocked(account, result.resetAt, "429")
|
||
if err := s.accountRepo.SetRateLimited(ctx, account.ID, result.resetAt); err != nil {
|
||
slog.Warn("rate_limit_set_failed", "account_id", account.ID, "error", err)
|
||
return
|
||
}
|
||
|
||
// 更新 session window:优先使用 5h-reset 头精确计算,否则从 resetAt 反推
|
||
windowEnd := result.resetAt
|
||
if result.fiveHourReset != nil {
|
||
windowEnd = *result.fiveHourReset
|
||
}
|
||
windowStart := windowEnd.Add(-5 * time.Hour)
|
||
if err := s.accountRepo.UpdateSessionWindow(ctx, account.ID, &windowStart, &windowEnd, "rejected"); err != nil {
|
||
slog.Warn("rate_limit_update_session_window_failed", "account_id", account.ID, "error", err)
|
||
}
|
||
|
||
slog.Info("anthropic_account_rate_limited", "account_id", account.ID, "reset_at", result.resetAt, "reset_in", time.Until(result.resetAt).Truncate(time.Second))
|
||
return
|
||
}
|
||
|
||
// 3. 尝试从响应头解析重置时间(Anthropic 聚合头,向后兼容)
|
||
resetTimestamp := headers.Get("anthropic-ratelimit-unified-reset")
|
||
|
||
// 4. 如果响应头没有,尝试从响应体解析(OpenAI usage_limit_reached, Gemini)
|
||
if resetTimestamp == "" {
|
||
switch account.Platform {
|
||
case PlatformOpenAI:
|
||
// 尝试解析 OpenAI 的 usage_limit_reached 错误
|
||
if resetAt := parseOpenAIRateLimitResetTime(responseBody); resetAt != nil {
|
||
resetTime := time.Unix(*resetAt, 0)
|
||
s.notifyAccountSchedulingBlocked(account, resetTime, "429")
|
||
if err := s.accountRepo.SetRateLimited(ctx, account.ID, resetTime); err != nil {
|
||
slog.Warn("rate_limit_set_failed", "account_id", account.ID, "error", err)
|
||
return
|
||
}
|
||
slog.Info("account_rate_limited", "account_id", account.ID, "platform", account.Platform, "reset_at", resetTime, "reset_in", time.Until(resetTime).Truncate(time.Second))
|
||
return
|
||
}
|
||
case PlatformGemini, PlatformAntigravity:
|
||
// 尝试解析 Gemini 格式(用于其他平台)
|
||
if resetAt := ParseGeminiRateLimitResetTime(responseBody); resetAt != nil {
|
||
resetTime := time.Unix(*resetAt, 0)
|
||
s.notifyAccountSchedulingBlocked(account, resetTime, "429")
|
||
if err := s.accountRepo.SetRateLimited(ctx, account.ID, resetTime); err != nil {
|
||
slog.Warn("rate_limit_set_failed", "account_id", account.ID, "error", err)
|
||
return
|
||
}
|
||
slog.Info("account_rate_limited", "account_id", account.ID, "platform", account.Platform, "reset_at", resetTime, "reset_in", time.Until(resetTime).Truncate(time.Second))
|
||
return
|
||
}
|
||
}
|
||
|
||
// Anthropic 平台:没有限流重置时间的 429 可能是非真实限流(如 Extra usage required),
|
||
// 不适合按 5h/7d 窗口长时间封禁;但完全不标记会导致账号永不冷却,
|
||
// 调度器让每个请求反复撞同一批持续 429 的账号(failover 预算被白白烧掉,
|
||
// 客户端稳定收到 429)。因此同样走可配置的秒级兜底回避,管理端可调大或关闭。
|
||
if account.Platform == PlatformAnthropic {
|
||
slog.Warn("rate_limit_429_no_reset_time",
|
||
"account_id", account.ID,
|
||
"platform", account.Platform,
|
||
"reason", "no rate limit reset time in headers, likely not a real rate limit")
|
||
s.apply429FallbackRateLimit(ctx, account, "anthropic_no_reset_time")
|
||
return
|
||
}
|
||
|
||
// 其他平台:没有重置时间,使用可配置的秒级默认回避,避免误伤长时间不可调度。
|
||
s.apply429FallbackRateLimit(ctx, account, "no_reset_time")
|
||
return
|
||
}
|
||
|
||
// 解析Unix时间戳
|
||
ts, err := strconv.ParseInt(resetTimestamp, 10, 64)
|
||
if err != nil {
|
||
slog.Warn("rate_limit_reset_parse_failed", "reset_timestamp", resetTimestamp, "error", err)
|
||
s.apply429FallbackRateLimit(ctx, account, "reset_parse_failed")
|
||
return
|
||
}
|
||
|
||
resetAt := time.Unix(ts, 0)
|
||
|
||
// 标记限流状态
|
||
s.notifyAccountSchedulingBlocked(account, resetAt, "429")
|
||
if err := s.accountRepo.SetRateLimited(ctx, account.ID, resetAt); err != nil {
|
||
slog.Warn("rate_limit_set_failed", "account_id", account.ID, "error", err)
|
||
return
|
||
}
|
||
|
||
// 根据重置时间反推5h窗口
|
||
windowEnd := resetAt
|
||
windowStart := resetAt.Add(-5 * time.Hour)
|
||
if err := s.accountRepo.UpdateSessionWindow(ctx, account.ID, &windowStart, &windowEnd, "rejected"); err != nil {
|
||
slog.Warn("rate_limit_update_session_window_failed", "account_id", account.ID, "error", err)
|
||
}
|
||
|
||
slog.Info("account_rate_limited", "account_id", account.ID, "reset_at", resetAt)
|
||
}
|
||
|
||
func (s *RateLimitService) apply429FallbackRateLimit(ctx context.Context, account *Account, reason string) {
|
||
cooldown, enabled := s.get429FallbackCooldown(ctx, account)
|
||
if !enabled {
|
||
slog.Info("rate_limit_429_fallback_ignored", "account_id", account.ID, "platform", account.Platform, "reason", reason)
|
||
return
|
||
}
|
||
|
||
resetAt := time.Now().Add(cooldown)
|
||
slog.Warn("rate_limit_429_fallback_used", "account_id", account.ID, "platform", account.Platform, "reason", reason, "using_default", cooldown.String())
|
||
s.notifyAccountSchedulingBlocked(account, resetAt, "429_fallback")
|
||
if err := s.accountRepo.SetRateLimited(ctx, account.ID, resetAt); err != nil {
|
||
slog.Warn("rate_limit_set_failed", "account_id", account.ID, "error", err)
|
||
}
|
||
}
|
||
|
||
func (s *RateLimitService) get429FallbackCooldown(ctx context.Context, account *Account) (time.Duration, bool) {
|
||
if s.settingService != nil {
|
||
settings, err := s.settingService.GetRateLimit429CooldownSettings(ctx)
|
||
if err == nil && settings != nil {
|
||
if !settings.Enabled {
|
||
return 0, false
|
||
}
|
||
seconds := clampRateLimit429CooldownSeconds(settings.CooldownSeconds)
|
||
return time.Duration(seconds) * time.Second, true
|
||
}
|
||
slog.Warn("rate_limit_429_settings_read_failed", "account_id", account.ID, "error", err)
|
||
}
|
||
|
||
seconds := defaultRateLimit429CooldownSeconds
|
||
seconds = clampRateLimit429CooldownSeconds(seconds)
|
||
return time.Duration(seconds) * time.Second, true
|
||
}
|
||
|
||
func clampRateLimit429CooldownSeconds(seconds int) int {
|
||
if seconds < 1 {
|
||
return 1
|
||
}
|
||
if seconds > maxRateLimit429CooldownSeconds {
|
||
return maxRateLimit429CooldownSeconds
|
||
}
|
||
return seconds
|
||
}
|
||
|
||
// calculateOpenAI429ResetTime 从 OpenAI 429 响应头计算正确的重置时间
|
||
// 返回 nil 表示无法从响应头中确定重置时间
|
||
func calculateOpenAI429ResetTime(headers http.Header) *time.Time {
|
||
snapshot := ParseCodexRateLimitHeaders(headers)
|
||
if snapshot == nil {
|
||
return nil
|
||
}
|
||
|
||
normalized := snapshot.Normalize()
|
||
if normalized == nil {
|
||
return nil
|
||
}
|
||
|
||
now := time.Now()
|
||
|
||
// 判断哪个限制被触发(used_percent >= 100)
|
||
is7dExhausted := normalized.Used7dPercent != nil && *normalized.Used7dPercent >= 100
|
||
is5hExhausted := normalized.Used5hPercent != nil && *normalized.Used5hPercent >= 100
|
||
|
||
// 优先使用被触发限制的重置时间
|
||
if is7dExhausted && normalized.Reset7dSeconds != nil {
|
||
resetAt := now.Add(time.Duration(*normalized.Reset7dSeconds) * time.Second)
|
||
slog.Info("openai_429_7d_limit_exhausted", "reset_after_seconds", *normalized.Reset7dSeconds, "reset_at", resetAt)
|
||
return &resetAt
|
||
}
|
||
if is5hExhausted && normalized.Reset5hSeconds != nil {
|
||
resetAt := now.Add(time.Duration(*normalized.Reset5hSeconds) * time.Second)
|
||
slog.Info("openai_429_5h_limit_exhausted", "reset_after_seconds", *normalized.Reset5hSeconds, "reset_at", resetAt)
|
||
return &resetAt
|
||
}
|
||
|
||
// 都未达到100%但收到429,使用较长的重置时间
|
||
var maxResetSecs int
|
||
if normalized.Reset7dSeconds != nil && *normalized.Reset7dSeconds > maxResetSecs {
|
||
maxResetSecs = *normalized.Reset7dSeconds
|
||
}
|
||
if normalized.Reset5hSeconds != nil && *normalized.Reset5hSeconds > maxResetSecs {
|
||
maxResetSecs = *normalized.Reset5hSeconds
|
||
}
|
||
if maxResetSecs > 0 {
|
||
resetAt := now.Add(time.Duration(maxResetSecs) * time.Second)
|
||
slog.Info("openai_429_using_max_reset", "max_reset_seconds", maxResetSecs, "reset_at", resetAt)
|
||
return &resetAt
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
func (s *RateLimitService) calculateOpenAI429ResetTime(headers http.Header) *time.Time {
|
||
return calculateOpenAI429ResetTime(headers)
|
||
}
|
||
|
||
// anthropic429Result holds the parsed Anthropic 429 rate-limit information.
|
||
type anthropic429Result struct {
|
||
resetAt time.Time // The correct reset time to use for SetRateLimited
|
||
fiveHourReset *time.Time // 5h window reset timestamp (for session window calculation), nil if not available
|
||
}
|
||
|
||
type anthropicWindowLimit struct {
|
||
window string
|
||
resetAt time.Time
|
||
reason string
|
||
}
|
||
|
||
func selectAnthropicExhaustedWindow(headers http.Header, now time.Time) *anthropicWindowLimit {
|
||
reset5h, ok5hReset := parseAnthropicWindowReset(headers, "5h", now)
|
||
reset7d, ok7dReset := parseAnthropicWindowReset(headers, "7d", now)
|
||
|
||
exceeded5h := isAnthropic5hRejected(headers) || isAnthropicWindowExceeded(headers, "5h")
|
||
exceeded7d := isAnthropicWindowExceeded(headers, "7d")
|
||
|
||
if exceeded7d && ok7dReset {
|
||
return &anthropicWindowLimit{
|
||
window: "7d",
|
||
resetAt: reset7d,
|
||
reason: "anthropic_7d_window_exhausted",
|
||
}
|
||
}
|
||
if exceeded5h && ok5hReset {
|
||
return &anthropicWindowLimit{
|
||
window: "5h",
|
||
resetAt: reset5h,
|
||
reason: "anthropic_5h_window_exhausted",
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func isAnthropic5hRejected(headers http.Header) bool {
|
||
return isAnthropicWindowRejected(headers, "5h")
|
||
}
|
||
|
||
func isAnthropicWindowRejected(headers http.Header, window string) bool {
|
||
return strings.EqualFold(strings.TrimSpace(headers.Get("anthropic-ratelimit-unified-"+window+"-status")), "rejected")
|
||
}
|
||
|
||
func parseAnthropicWindowReset(headers http.Header, window string, now time.Time) (time.Time, bool) {
|
||
maxAge := 8 * 24 * time.Hour
|
||
if window == "5h" {
|
||
maxAge = 6 * time.Hour
|
||
}
|
||
return parseAnthropicResetTimestamp(headers.Get("anthropic-ratelimit-unified-"+window+"-reset"), now, maxAge)
|
||
}
|
||
|
||
// parseAnthropicResetTimestamp 解析 Anthropic reset 头的 Unix 时间戳(自动识别毫秒),
|
||
// 并校验落在 (now, now+maxAge] 的合理区间内。
|
||
func parseAnthropicResetTimestamp(raw string, now time.Time, maxAge time.Duration) (time.Time, bool) {
|
||
raw = strings.TrimSpace(raw)
|
||
if raw == "" {
|
||
return time.Time{}, false
|
||
}
|
||
ts, err := strconv.ParseInt(raw, 10, 64)
|
||
if err != nil {
|
||
return time.Time{}, false
|
||
}
|
||
if ts > 1e11 {
|
||
ts = ts / 1000
|
||
}
|
||
resetAt := time.Unix(ts, 0)
|
||
if !resetAt.After(now) || resetAt.After(now.Add(maxAge)) {
|
||
return time.Time{}, false
|
||
}
|
||
return resetAt, true
|
||
}
|
||
|
||
func shouldPersistAnthropicWindowLimit(account *Account, limit *anthropicWindowLimit, now time.Time) bool {
|
||
if account == nil || limit == nil || !limit.resetAt.After(now) {
|
||
return false
|
||
}
|
||
if account.RateLimitResetAt == nil {
|
||
return true
|
||
}
|
||
if !account.RateLimitResetAt.After(now) {
|
||
return true
|
||
}
|
||
return limit.resetAt.After(*account.RateLimitResetAt)
|
||
}
|
||
|
||
func (s *RateLimitService) persistAnthropicExhaustedWindowLimit(ctx context.Context, account *Account, headers http.Header) bool {
|
||
if s == nil || s.accountRepo == nil || account == nil {
|
||
return false
|
||
}
|
||
now := time.Now()
|
||
limit := selectAnthropicExhaustedWindow(headers, now)
|
||
if limit == nil {
|
||
return false
|
||
}
|
||
if !shouldPersistAnthropicWindowLimit(account, limit, now) {
|
||
slog.Info("anthropic_window_rate_limit_kept",
|
||
"account_id", account.ID,
|
||
"window", limit.window,
|
||
"reset_at", limit.resetAt,
|
||
"existing_reset_at", account.RateLimitResetAt)
|
||
return true
|
||
}
|
||
|
||
s.notifyAccountSchedulingBlocked(account, limit.resetAt, limit.reason)
|
||
if err := s.accountRepo.SetRateLimited(ctx, account.ID, limit.resetAt); err != nil {
|
||
slog.Warn("anthropic_window_rate_limit_set_failed",
|
||
"account_id", account.ID,
|
||
"window", limit.window,
|
||
"reset_at", limit.resetAt,
|
||
"error", err)
|
||
return true
|
||
}
|
||
slog.Info("anthropic_window_rate_limited",
|
||
"account_id", account.ID,
|
||
"window", limit.window,
|
||
"reset_at", limit.resetAt,
|
||
"reset_in", time.Until(limit.resetAt).Truncate(time.Second))
|
||
return true
|
||
}
|
||
|
||
const anthropicFableWindowReason = "anthropic_7d_oi_window_exhausted"
|
||
|
||
// selectAnthropicFableWindowLimit parses the Anthropic 7d_oi per-model window
|
||
// headers (the Fable-only 7d window, e.g. anthropic-ratelimit-unified-7d_oi-*).
|
||
// Unlike 5h/7d, exhaustion of this window only limits the Fable model family —
|
||
// the account must stay schedulable for other models.
|
||
//
|
||
// The 7d_oi surpassed-threshold header carries a float ("1.0") rather than
|
||
// "true", so exhaustion is detected via status=rejected or utilization >= 1.0.
|
||
// When the 7d_oi reset header is missing, the aggregated
|
||
// anthropic-ratelimit-unified-reset is used (it mirrors the binding claim's
|
||
// reset when 7d_oi is the representative claim).
|
||
func selectAnthropicFableWindowLimit(headers http.Header, now time.Time) *anthropicWindowLimit {
|
||
if !isAnthropicWindowRejected(headers, "7d_oi") && !isAnthropicWindowExceeded(headers, "7d_oi") {
|
||
return nil
|
||
}
|
||
resetAt, ok := parseAnthropicWindowReset(headers, "7d_oi", now)
|
||
if !ok {
|
||
resetAt, ok = parseAnthropicAggregateReset(headers, now)
|
||
}
|
||
if !ok {
|
||
return nil
|
||
}
|
||
return &anthropicWindowLimit{
|
||
window: "7d_oi",
|
||
resetAt: resetAt,
|
||
reason: anthropicFableWindowReason,
|
||
}
|
||
}
|
||
|
||
// parseAnthropicAggregateReset parses the aggregated
|
||
// anthropic-ratelimit-unified-reset header with the same sanity checks as the
|
||
// per-window variant (7d scale).
|
||
func parseAnthropicAggregateReset(headers http.Header, now time.Time) (time.Time, bool) {
|
||
return parseAnthropicResetTimestamp(headers.Get("anthropic-ratelimit-unified-reset"), now, 8*24*time.Hour)
|
||
}
|
||
|
||
// persistAnthropicFableWindowLimit marks the Fable model family as rate limited
|
||
// when the 7d_oi window is exhausted. Returns true when the 7d_oi window was the
|
||
// (or a) trigger of this 429, so the caller must not fall through to logic that
|
||
// would mark the whole account as rate limited.
|
||
func (s *RateLimitService) persistAnthropicFableWindowLimit(ctx context.Context, account *Account, headers http.Header) bool {
|
||
if s == nil || s.accountRepo == nil || account == nil {
|
||
return false
|
||
}
|
||
now := time.Now()
|
||
limit := selectAnthropicFableWindowLimit(headers, now)
|
||
if limit == nil {
|
||
return false
|
||
}
|
||
// 429 响应头本身携带最新的窗口用量(7d_oi utilization=1.0)。限流期内
|
||
// Fable 请求不再调度到该账号,若不在此处采样,7d F 进度条会冻结在
|
||
// 限流前的旧值直到窗口重置。
|
||
s.samplePassiveUsageFromHeaders(ctx, account, headers)
|
||
if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, anthropicFableRateLimitKey, limit.resetAt, limit.reason); err != nil {
|
||
slog.Warn("anthropic_fable_window_rate_limit_set_failed",
|
||
"account_id", account.ID,
|
||
"scope", anthropicFableRateLimitKey,
|
||
"reset_at", limit.resetAt,
|
||
"error", err)
|
||
return true
|
||
}
|
||
slog.Info("anthropic_fable_window_model_rate_limited",
|
||
"account_id", account.ID,
|
||
"scope", anthropicFableRateLimitKey,
|
||
"reset_at", limit.resetAt,
|
||
"reset_in", time.Until(limit.resetAt).Truncate(time.Second))
|
||
return true
|
||
}
|
||
|
||
// calculateAnthropic429ResetTime parses Anthropic's per-window rate-limit headers
|
||
// to determine which window (5h or 7d) actually triggered the 429.
|
||
//
|
||
// Headers used:
|
||
// - anthropic-ratelimit-unified-5h-utilization / anthropic-ratelimit-unified-5h-surpassed-threshold
|
||
// - anthropic-ratelimit-unified-5h-reset
|
||
// - anthropic-ratelimit-unified-7d-utilization / anthropic-ratelimit-unified-7d-surpassed-threshold
|
||
// - anthropic-ratelimit-unified-7d-reset
|
||
//
|
||
// Returns nil when the per-window headers are absent (caller should fall back to
|
||
// the aggregated anthropic-ratelimit-unified-reset header).
|
||
func calculateAnthropic429ResetTime(headers http.Header) *anthropic429Result {
|
||
reset5hStr := headers.Get("anthropic-ratelimit-unified-5h-reset")
|
||
reset7dStr := headers.Get("anthropic-ratelimit-unified-7d-reset")
|
||
|
||
if reset5hStr == "" && reset7dStr == "" {
|
||
return nil
|
||
}
|
||
|
||
var reset5h, reset7d *time.Time
|
||
if ts, err := strconv.ParseInt(reset5hStr, 10, 64); err == nil {
|
||
t := time.Unix(ts, 0)
|
||
reset5h = &t
|
||
}
|
||
if ts, err := strconv.ParseInt(reset7dStr, 10, 64); err == nil {
|
||
t := time.Unix(ts, 0)
|
||
reset7d = &t
|
||
}
|
||
|
||
is5hExceeded := isAnthropicWindowExceeded(headers, "5h")
|
||
is7dExceeded := isAnthropicWindowExceeded(headers, "7d")
|
||
|
||
slog.Info("anthropic_429_window_analysis",
|
||
"is_5h_exceeded", is5hExceeded,
|
||
"is_7d_exceeded", is7dExceeded,
|
||
"reset_5h", reset5hStr,
|
||
"reset_7d", reset7dStr,
|
||
)
|
||
|
||
// Select the correct reset time based on which window(s) are exceeded.
|
||
var chosen *time.Time
|
||
switch {
|
||
case is5hExceeded && is7dExceeded:
|
||
// Both exceeded → prefer 7d (longer cooldown), fall back to 5h
|
||
chosen = reset7d
|
||
if chosen == nil {
|
||
chosen = reset5h
|
||
}
|
||
case is5hExceeded:
|
||
chosen = reset5h
|
||
case is7dExceeded:
|
||
chosen = reset7d
|
||
default:
|
||
// Neither flag clearly exceeded — pick the sooner reset as best guess
|
||
chosen = pickSooner(reset5h, reset7d)
|
||
}
|
||
|
||
if chosen == nil {
|
||
return nil
|
||
}
|
||
return &anthropic429Result{resetAt: *chosen, fiveHourReset: reset5h}
|
||
}
|
||
|
||
// isAnthropicWindowExceeded checks whether a given Anthropic rate-limit window
|
||
// (e.g. "5h" or "7d") has been exceeded, using utilization and surpassed-threshold headers.
|
||
func isAnthropicWindowExceeded(headers http.Header, window string) bool {
|
||
prefix := "anthropic-ratelimit-unified-" + window + "-"
|
||
|
||
// Check surpassed-threshold first (most explicit signal)
|
||
if st := headers.Get(prefix + "surpassed-threshold"); strings.EqualFold(st, "true") {
|
||
return true
|
||
}
|
||
|
||
// Fall back to utilization >= 1.0
|
||
if utilStr := headers.Get(prefix + "utilization"); utilStr != "" {
|
||
if util, err := strconv.ParseFloat(utilStr, 64); err == nil && util >= 1.0-1e-9 {
|
||
// Use a small epsilon to handle floating point: treat 0.9999999... as >= 1.0
|
||
return true
|
||
}
|
||
}
|
||
|
||
return false
|
||
}
|
||
|
||
// pickSooner returns whichever of the two time pointers is earlier.
|
||
// If only one is non-nil, it is returned. If both are nil, returns nil.
|
||
func pickSooner(a, b *time.Time) *time.Time {
|
||
switch {
|
||
case a != nil && b != nil:
|
||
if a.Before(*b) {
|
||
return a
|
||
}
|
||
return b
|
||
case a != nil:
|
||
return a
|
||
default:
|
||
return b
|
||
}
|
||
}
|
||
|
||
func (s *RateLimitService) persistOpenAICodexSnapshot(ctx context.Context, account *Account, headers http.Header) {
|
||
if s == nil || s.accountRepo == nil || account == nil || headers == nil {
|
||
return
|
||
}
|
||
// spark 影子的 codex_* 仅由 QueryUsage(/wham/usage bengalfox 道)更新,不能被 /responses 的
|
||
// x-codex-* 全局头快照污染(外审第7轮 P1,与 updateCodexUsageSnapshot 同口径)。
|
||
if account.IsShadow() {
|
||
return
|
||
}
|
||
snapshot := ParseCodexRateLimitHeaders(headers)
|
||
if snapshot == nil {
|
||
return
|
||
}
|
||
updates := buildCodexUsageExtraUpdates(snapshot, time.Now())
|
||
if len(updates) == 0 {
|
||
return
|
||
}
|
||
if err := s.accountRepo.UpdateExtra(ctx, account.ID, updates); err != nil {
|
||
slog.Warn("openai_codex_snapshot_persist_failed", "account_id", account.ID, "error", err)
|
||
}
|
||
}
|
||
|
||
// parseOpenAIRateLimitResetTime 解析 OpenAI 格式的 429 响应,返回重置时间的 Unix 时间戳
|
||
// OpenAI 的 usage_limit_reached 错误格式:
|
||
//
|
||
// {
|
||
// "error": {
|
||
// "message": "The usage limit has been reached",
|
||
// "type": "usage_limit_reached",
|
||
// "resets_at": 1769404154,
|
||
// "resets_in_seconds": 133107
|
||
// }
|
||
// }
|
||
func parseOpenAIRateLimitResetTime(body []byte) *int64 {
|
||
var parsed map[string]any
|
||
if err := json.Unmarshal(body, &parsed); err != nil {
|
||
return nil
|
||
}
|
||
|
||
errObj, ok := parsed["error"].(map[string]any)
|
||
if !ok {
|
||
return nil
|
||
}
|
||
|
||
// 检查是否为 usage_limit_reached 或 rate_limit_exceeded 类型
|
||
errType, _ := errObj["type"].(string)
|
||
if errType != "usage_limit_reached" && errType != "rate_limit_exceeded" {
|
||
return nil
|
||
}
|
||
|
||
// 优先使用 resets_at(Unix 时间戳)
|
||
if resetsAt, ok := errObj["resets_at"].(float64); ok {
|
||
ts := int64(resetsAt)
|
||
return &ts
|
||
}
|
||
if resetsAt, ok := errObj["resets_at"].(string); ok {
|
||
if ts, err := strconv.ParseInt(resetsAt, 10, 64); err == nil {
|
||
return &ts
|
||
}
|
||
}
|
||
|
||
// 如果没有 resets_at,尝试使用 resets_in_seconds
|
||
if resetsInSeconds, ok := errObj["resets_in_seconds"].(float64); ok {
|
||
ts := time.Now().Unix() + int64(resetsInSeconds)
|
||
return &ts
|
||
}
|
||
if resetsInSeconds, ok := errObj["resets_in_seconds"].(string); ok {
|
||
if sec, err := strconv.ParseInt(resetsInSeconds, 10, 64); err == nil {
|
||
ts := time.Now().Unix() + sec
|
||
return &ts
|
||
}
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
func parseOpenAIRateLimitPlanType(body []byte) string {
|
||
var parsed map[string]any
|
||
if err := json.Unmarshal(body, &parsed); err != nil {
|
||
return ""
|
||
}
|
||
|
||
errObj, ok := parsed["error"].(map[string]any)
|
||
if !ok {
|
||
return ""
|
||
}
|
||
|
||
errType, _ := errObj["type"].(string)
|
||
if errType != "usage_limit_reached" && errType != "rate_limit_exceeded" {
|
||
return ""
|
||
}
|
||
|
||
planType, _ := errObj["plan_type"].(string)
|
||
return strings.ToLower(strings.TrimSpace(planType))
|
||
}
|
||
|
||
func persistOpenAI429PlanType(ctx context.Context, repo AccountRepository, account *Account, body []byte) {
|
||
if repo == nil || account == nil || account.Platform != PlatformOpenAI {
|
||
return
|
||
}
|
||
// spark 影子账号恒不持凭据:即便收到带 plan_type 的 429,也不能把 plan_type 写进影子 credentials
|
||
// ——该路径走 repo.BulkUpdate 直写、不经 persistAccountCredentials 守卫(外审第7轮 P1)。
|
||
// plan_type 由母账号在自己的请求上维护,影子跳过。
|
||
if account.IsCredentialShadow() {
|
||
return
|
||
}
|
||
|
||
planType := parseOpenAIRateLimitPlanType(body)
|
||
if planType == "" {
|
||
return
|
||
}
|
||
|
||
current := strings.TrimSpace(account.GetCredential("plan_type"))
|
||
if strings.EqualFold(current, planType) {
|
||
return
|
||
}
|
||
|
||
if _, err := repo.BulkUpdate(ctx, []int64{account.ID}, AccountBulkUpdate{
|
||
Credentials: map[string]any{"plan_type": planType},
|
||
}); err != nil {
|
||
slog.Warn("openai_429_plan_type_sync_failed", "account_id", account.ID, "plan_type", planType, "error", err)
|
||
return
|
||
}
|
||
|
||
if account.Credentials == nil {
|
||
account.Credentials = make(map[string]any, 1)
|
||
}
|
||
account.Credentials["plan_type"] = planType
|
||
slog.Info("openai_429_plan_type_synced", "account_id", account.ID, "previous_plan_type", current, "plan_type", planType)
|
||
}
|
||
|
||
// handle529 处理529过载错误
|
||
// 根据配置决定是否暂停账号调度及冷却时长
|
||
func (s *RateLimitService) handle529(ctx context.Context, account *Account) {
|
||
var settings *OverloadCooldownSettings
|
||
if s.settingService != nil {
|
||
var err error
|
||
settings, err = s.settingService.GetOverloadCooldownSettings(ctx)
|
||
if err != nil {
|
||
slog.Warn("overload_settings_read_failed", "account_id", account.ID, "error", err)
|
||
settings = nil
|
||
}
|
||
}
|
||
// 回退到配置文件
|
||
if settings == nil {
|
||
cooldown := s.cfg.RateLimit.OverloadCooldownMinutes
|
||
if cooldown <= 0 {
|
||
cooldown = 10
|
||
}
|
||
settings = &OverloadCooldownSettings{Enabled: true, CooldownMinutes: cooldown}
|
||
}
|
||
|
||
if !settings.Enabled {
|
||
slog.Info("account_529_ignored", "account_id", account.ID, "reason", "overload_cooldown_disabled")
|
||
return
|
||
}
|
||
|
||
cooldownMinutes := settings.CooldownMinutes
|
||
if cooldownMinutes <= 0 {
|
||
cooldownMinutes = 10
|
||
}
|
||
|
||
until := time.Now().Add(time.Duration(cooldownMinutes) * time.Minute)
|
||
s.notifyAccountSchedulingBlocked(account, until, "529")
|
||
if err := s.accountRepo.SetOverloaded(ctx, account.ID, until); err != nil {
|
||
slog.Warn("overload_set_failed", "account_id", account.ID, "error", err)
|
||
return
|
||
}
|
||
|
||
slog.Info("account_overloaded", "account_id", account.ID, "until", until)
|
||
}
|
||
|
||
// UpdateSessionWindow 从成功响应更新5h窗口状态
|
||
func (s *RateLimitService) UpdateSessionWindow(ctx context.Context, account *Account, headers http.Header) {
|
||
status := headers.Get("anthropic-ratelimit-unified-5h-status")
|
||
if status == "" {
|
||
return
|
||
}
|
||
|
||
// 检查是否需要初始化时间窗口
|
||
// 对于 Setup Token 账号,首次成功请求时需要预测时间窗口
|
||
var windowStart, windowEnd *time.Time
|
||
needInitWindow := account.SessionWindowEnd == nil || time.Now().After(*account.SessionWindowEnd)
|
||
|
||
// 优先使用响应头中的真实重置时间(比预测更准确)
|
||
if resetStr := headers.Get("anthropic-ratelimit-unified-5h-reset"); resetStr != "" {
|
||
if ts, err := strconv.ParseInt(resetStr, 10, 64); err == nil {
|
||
// 检测可能的毫秒时间戳(秒级约为 1e9,毫秒约为 1e12)
|
||
if ts > 1e11 {
|
||
slog.Warn("account_session_window_header_millis_detected", "account_id", account.ID, "raw_reset", resetStr)
|
||
ts = ts / 1000
|
||
}
|
||
end := time.Unix(ts, 0)
|
||
// 校验时间戳是否在合理范围内(不早于 5h 前,不晚于 7 天后)
|
||
minAllowed := time.Now().Add(-5 * time.Hour)
|
||
maxAllowed := time.Now().Add(7 * 24 * time.Hour)
|
||
if end.Before(minAllowed) || end.After(maxAllowed) {
|
||
slog.Warn("account_session_window_header_out_of_range", "account_id", account.ID, "raw_reset", resetStr, "parsed_end", end)
|
||
} else if needInitWindow || account.SessionWindowEnd == nil || !end.Equal(*account.SessionWindowEnd) {
|
||
// 窗口需要初始化,或者真实重置时间与已存储的不同,则更新
|
||
start := end.Add(-5 * time.Hour)
|
||
windowStart = &start
|
||
windowEnd = &end
|
||
slog.Info("account_session_window_from_header", "account_id", account.ID, "window_start", start, "window_end", end, "status", status)
|
||
}
|
||
} else {
|
||
slog.Warn("account_session_window_header_parse_failed", "account_id", account.ID, "raw_reset", resetStr, "error", err)
|
||
}
|
||
}
|
||
|
||
// 回退:如果没有真实重置时间且需要初始化窗口,使用预测
|
||
if windowEnd == nil && needInitWindow && (status == "allowed" || status == "allowed_warning") {
|
||
now := time.Now()
|
||
start := time.Date(now.Year(), now.Month(), now.Day(), now.Hour(), 0, 0, 0, now.Location())
|
||
end := start.Add(5 * time.Hour)
|
||
windowStart = &start
|
||
windowEnd = &end
|
||
slog.Info("account_session_window_initialized", "account_id", account.ID, "window_start", start, "window_end", end, "status", status)
|
||
}
|
||
|
||
// 窗口重置时清除旧的 utilization 和被动采样数据,避免残留上个窗口的数据
|
||
if windowEnd != nil && needInitWindow {
|
||
_ = s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{
|
||
"session_window_utilization": nil,
|
||
"passive_usage_7d_utilization": nil,
|
||
"passive_usage_7d_reset": nil,
|
||
"passive_usage_7d_oi_utilization": nil,
|
||
"passive_usage_7d_oi_reset": nil,
|
||
"passive_usage_sampled_at": nil,
|
||
})
|
||
}
|
||
|
||
if err := s.accountRepo.UpdateSessionWindow(ctx, account.ID, windowStart, windowEnd, status); err != nil {
|
||
slog.Warn("session_window_update_failed", "account_id", account.ID, "error", err)
|
||
}
|
||
|
||
// 被动采样:从响应头收集 5h + 7d + 7d_oi utilization,合并为一次 DB 写入
|
||
s.samplePassiveUsageFromHeaders(ctx, account, headers)
|
||
|
||
// 如果状态为allowed且之前有限流,说明窗口已重置,清除限流状态
|
||
if status == "allowed" && account.IsRateLimited() {
|
||
if err := s.ClearRateLimit(ctx, account.ID); err != nil {
|
||
slog.Warn("rate_limit_clear_failed", "account_id", account.ID, "error", err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// samplePassiveUsageFromHeaders 从 Anthropic 响应头收集 5h/7d/7d_oi 的
|
||
// utilization 与 reset 被动采样数据,合并为一次 Extra 写入。无数据时不写。
|
||
func (s *RateLimitService) samplePassiveUsageFromHeaders(ctx context.Context, account *Account, headers http.Header) {
|
||
extraUpdates := make(map[string]any, 6)
|
||
// 5h utilization(0-1 小数),供 estimateSetupTokenUsage 使用
|
||
if utilStr := headers.Get("anthropic-ratelimit-unified-5h-utilization"); utilStr != "" {
|
||
if util, err := strconv.ParseFloat(utilStr, 64); err == nil {
|
||
extraUpdates["session_window_utilization"] = util
|
||
}
|
||
}
|
||
// 7d utilization(0-1 小数)
|
||
if utilStr := headers.Get("anthropic-ratelimit-unified-7d-utilization"); utilStr != "" {
|
||
if util, err := strconv.ParseFloat(utilStr, 64); err == nil {
|
||
extraUpdates["passive_usage_7d_utilization"] = util
|
||
}
|
||
}
|
||
// 7d reset timestamp
|
||
if resetStr := headers.Get("anthropic-ratelimit-unified-7d-reset"); resetStr != "" {
|
||
if ts, err := strconv.ParseInt(resetStr, 10, 64); err == nil {
|
||
if ts > 1e11 {
|
||
ts = ts / 1000
|
||
}
|
||
extraUpdates["passive_usage_7d_reset"] = ts
|
||
}
|
||
}
|
||
// 7d_oi (Fable 专属 7d 窗口) utilization(0-1 小数)
|
||
if utilStr := headers.Get("anthropic-ratelimit-unified-7d_oi-utilization"); utilStr != "" {
|
||
if util, err := strconv.ParseFloat(utilStr, 64); err == nil {
|
||
extraUpdates["passive_usage_7d_oi_utilization"] = util
|
||
}
|
||
}
|
||
// 7d_oi reset timestamp
|
||
if resetStr := headers.Get("anthropic-ratelimit-unified-7d_oi-reset"); resetStr != "" {
|
||
if ts, err := strconv.ParseInt(resetStr, 10, 64); err == nil {
|
||
if ts > 1e11 {
|
||
ts = ts / 1000
|
||
}
|
||
extraUpdates["passive_usage_7d_oi_reset"] = ts
|
||
}
|
||
}
|
||
if len(extraUpdates) > 0 {
|
||
extraUpdates["passive_usage_sampled_at"] = time.Now().UTC().Format(time.RFC3339)
|
||
if err := s.accountRepo.UpdateExtra(ctx, account.ID, extraUpdates); err != nil {
|
||
slog.Warn("passive_usage_update_failed", "account_id", account.ID, "error", err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// ClearRateLimit 清除账号的限流状态
|
||
func (s *RateLimitService) ClearRateLimit(ctx context.Context, accountID int64) error {
|
||
if err := s.accountRepo.ClearRateLimit(ctx, accountID); err != nil {
|
||
return err
|
||
}
|
||
if err := s.accountRepo.ClearAntigravityQuotaScopes(ctx, accountID); err != nil {
|
||
return err
|
||
}
|
||
if err := s.accountRepo.ClearModelRateLimits(ctx, accountID); err != nil {
|
||
return err
|
||
}
|
||
// 清除限流时一并清理临时不可调度状态,避免周限/窗口重置后仍被本地临时状态阻断。
|
||
if err := s.accountRepo.ClearTempUnschedulable(ctx, accountID); err != nil {
|
||
return err
|
||
}
|
||
if s.tempUnschedCache != nil {
|
||
if err := s.tempUnschedCache.DeleteTempUnsched(ctx, accountID); err != nil {
|
||
slog.Warn("temp_unsched_cache_delete_failed", "account_id", accountID, "error", err)
|
||
}
|
||
}
|
||
s.ResetOpenAI403Counter(ctx, accountID)
|
||
s.notifyAccountSchedulingBlockCleared(accountID)
|
||
return nil
|
||
}
|
||
|
||
func (s *RateLimitService) ResetOpenAI403Counter(ctx context.Context, accountID int64) {
|
||
if s == nil || s.openAI403CounterCache == nil || accountID <= 0 {
|
||
return
|
||
}
|
||
if err := s.openAI403CounterCache.ResetOpenAI403Count(ctx, accountID); err != nil {
|
||
slog.Warn("openai_403_reset_failed", "account_id", accountID, "error", err)
|
||
}
|
||
}
|
||
|
||
// RecoverAccountState 按需恢复账号的可恢复运行时状态。
|
||
func (s *RateLimitService) RecoverAccountState(ctx context.Context, accountID int64, options AccountRecoveryOptions) (*SuccessfulTestRecoveryResult, error) {
|
||
account, err := s.accountRepo.GetByID(ctx, accountID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
result := &SuccessfulTestRecoveryResult{}
|
||
if account.Status == StatusError {
|
||
if err := s.accountRepo.ClearError(ctx, accountID); err != nil {
|
||
return nil, err
|
||
}
|
||
result.ClearedError = true
|
||
if options.InvalidateToken && s.tokenCacheInvalidator != nil && account.IsOAuth() {
|
||
if invalidateErr := s.tokenCacheInvalidator.InvalidateToken(ctx, account); invalidateErr != nil {
|
||
slog.Warn("recover_account_state_invalidate_token_failed", "account_id", accountID, "error", invalidateErr)
|
||
}
|
||
}
|
||
}
|
||
|
||
if hasRecoverableRuntimeState(account) {
|
||
if err := s.ClearRateLimit(ctx, accountID); err != nil {
|
||
return nil, err
|
||
}
|
||
result.ClearedRateLimit = true
|
||
}
|
||
if result.ClearedError || result.ClearedRateLimit {
|
||
s.ResetOpenAI403Counter(ctx, accountID)
|
||
if result.ClearedError && !result.ClearedRateLimit {
|
||
s.notifyAccountSchedulingBlockCleared(accountID)
|
||
}
|
||
}
|
||
|
||
return result, nil
|
||
}
|
||
|
||
// RecoverAccountAfterSuccessfulTest 将一次成功测试视为正常请求,
|
||
// 按需恢复 error / rate-limit / overload / temp-unsched / model-rate-limit 等运行时状态。
|
||
func (s *RateLimitService) RecoverAccountAfterSuccessfulTest(ctx context.Context, accountID int64) (*SuccessfulTestRecoveryResult, error) {
|
||
return s.RecoverAccountState(ctx, accountID, AccountRecoveryOptions{})
|
||
}
|
||
|
||
func (s *RateLimitService) ClearTempUnschedulable(ctx context.Context, accountID int64) error {
|
||
if err := s.accountRepo.ClearTempUnschedulable(ctx, accountID); err != nil {
|
||
return err
|
||
}
|
||
if s.tempUnschedCache != nil {
|
||
if err := s.tempUnschedCache.DeleteTempUnsched(ctx, accountID); err != nil {
|
||
slog.Warn("temp_unsched_cache_delete_failed", "account_id", accountID, "error", err)
|
||
}
|
||
}
|
||
// 同时清除模型级别限流
|
||
if err := s.accountRepo.ClearModelRateLimits(ctx, accountID); err != nil {
|
||
slog.Warn("clear_model_rate_limits_on_temp_unsched_reset_failed", "account_id", accountID, "error", err)
|
||
}
|
||
s.notifyAccountSchedulingBlockCleared(accountID)
|
||
return nil
|
||
}
|
||
|
||
func hasRecoverableRuntimeState(account *Account) bool {
|
||
if account == nil {
|
||
return false
|
||
}
|
||
if account.RateLimitedAt != nil || account.RateLimitResetAt != nil || account.OverloadUntil != nil || account.TempUnschedulableUntil != nil {
|
||
return true
|
||
}
|
||
if len(account.Extra) == 0 {
|
||
return false
|
||
}
|
||
return hasNonEmptyMapValue(account.Extra, "model_rate_limits") ||
|
||
hasNonEmptyMapValue(account.Extra, "antigravity_quota_scopes")
|
||
}
|
||
|
||
func hasNonEmptyMapValue(extra map[string]any, key string) bool {
|
||
raw, ok := extra[key]
|
||
if !ok || raw == nil {
|
||
return false
|
||
}
|
||
switch typed := raw.(type) {
|
||
case map[string]any:
|
||
return len(typed) > 0
|
||
case map[string]string:
|
||
return len(typed) > 0
|
||
case []any:
|
||
return len(typed) > 0
|
||
default:
|
||
return true
|
||
}
|
||
}
|
||
|
||
func (s *RateLimitService) GetTempUnschedStatus(ctx context.Context, accountID int64) (*TempUnschedState, error) {
|
||
now := time.Now().Unix()
|
||
if s.tempUnschedCache != nil {
|
||
state, err := s.tempUnschedCache.GetTempUnsched(ctx, accountID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if state != nil && state.UntilUnix > now {
|
||
return state, nil
|
||
}
|
||
}
|
||
|
||
account, err := s.accountRepo.GetByID(ctx, accountID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if account.TempUnschedulableUntil == nil {
|
||
return nil, nil
|
||
}
|
||
if account.TempUnschedulableUntil.Unix() <= now {
|
||
return nil, nil
|
||
}
|
||
|
||
state := &TempUnschedState{
|
||
UntilUnix: account.TempUnschedulableUntil.Unix(),
|
||
}
|
||
|
||
if account.TempUnschedulableReason != "" {
|
||
var parsed TempUnschedState
|
||
if err := json.Unmarshal([]byte(account.TempUnschedulableReason), &parsed); err == nil {
|
||
if parsed.UntilUnix == 0 {
|
||
parsed.UntilUnix = state.UntilUnix
|
||
}
|
||
state = &parsed
|
||
} else {
|
||
state.ErrorMessage = account.TempUnschedulableReason
|
||
}
|
||
}
|
||
|
||
if s.tempUnschedCache != nil {
|
||
if err := s.tempUnschedCache.SetTempUnsched(ctx, accountID, state); err != nil {
|
||
slog.Warn("temp_unsched_cache_set_failed", "account_id", accountID, "error", err)
|
||
}
|
||
}
|
||
|
||
return state, nil
|
||
}
|
||
|
||
func (s *RateLimitService) HandleTempUnschedulable(ctx context.Context, account *Account, statusCode int, responseBody []byte, requestedModel ...string) bool {
|
||
if account == nil {
|
||
return false
|
||
}
|
||
if account.IsPoolMode() && !account.IsCustomErrorCodesEnabled() {
|
||
return false
|
||
}
|
||
if !account.ShouldHandleErrorCode(statusCode) {
|
||
return false
|
||
}
|
||
ctx = withTempUnschedulableModel(ctx, requestedModel)
|
||
return s.tryTempUnschedulable(ctx, account, statusCode, responseBody, firstRequestedModel(requestedModel))
|
||
}
|
||
|
||
func (s *RateLimitService) HandleOpenAIImageRateLimit(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte) bool {
|
||
if s == nil || account == nil || s.accountRepo == nil {
|
||
return false
|
||
}
|
||
if account.Platform != PlatformOpenAI {
|
||
return false
|
||
}
|
||
if !account.ShouldHandleErrorCode(statusCode) {
|
||
slog.Info("openai_image_rate_limit_skipped_by_error_code_policy", "account_id", account.ID, "status_code", statusCode)
|
||
return false
|
||
}
|
||
if !isOpenAIImageRateLimitError(statusCode, responseBody) {
|
||
return false
|
||
}
|
||
|
||
resetAt := openAIImageRateLimitResetAt(headers, responseBody)
|
||
if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, openAIImageGenerationRateLimitKey, resetAt, openAIImageRateLimitReason); err != nil {
|
||
slog.Warn("openai_image_rate_limit_set_model_rate_limit_failed", "account_id", account.ID, "scope", openAIImageGenerationRateLimitKey, "error", err)
|
||
return true
|
||
}
|
||
slog.Info("openai_image_rate_limited", "account_id", account.ID, "scope", openAIImageGenerationRateLimitKey, "reset_at", resetAt, "reset_in", time.Until(resetAt).Truncate(time.Second))
|
||
return true
|
||
}
|
||
|
||
func isOpenAIImageRateLimitError(statusCode int, body []byte) bool {
|
||
if statusCode != http.StatusTooManyRequests || len(body) == 0 {
|
||
return false
|
||
}
|
||
lower := strings.ToLower(string(body))
|
||
for _, marker := range []string{
|
||
"for limit gpt-image",
|
||
"input-images per min",
|
||
"gpt-image-2-codex",
|
||
"gpt-image",
|
||
} {
|
||
if strings.Contains(lower, marker) {
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}
|
||
|
||
func openAIImageRateLimitResetAt(headers http.Header, body []byte) time.Time {
|
||
now := time.Now()
|
||
if resetAt := parseRetryAfterResetTime(headers, now); resetAt != nil && resetAt.After(now) {
|
||
return *resetAt
|
||
}
|
||
if resetAt := calculateOpenAI429ResetTime(headers); resetAt != nil && resetAt.After(now) {
|
||
return *resetAt
|
||
}
|
||
if resetUnix := parseOpenAIRateLimitResetTime(body); resetUnix != nil {
|
||
if resetAt := time.Unix(*resetUnix, 0); resetAt.After(now) {
|
||
return resetAt
|
||
}
|
||
}
|
||
if cooldown := parseOpenAIImageTryAgainCooldown(body); cooldown > 0 {
|
||
return now.Add(cooldown)
|
||
}
|
||
return now.Add(openAIImageRateLimitDefaultCooldown)
|
||
}
|
||
|
||
func parseRetryAfterResetTime(headers http.Header, now time.Time) *time.Time {
|
||
if headers == nil {
|
||
return nil
|
||
}
|
||
raw := strings.TrimSpace(headers.Get("Retry-After"))
|
||
if raw == "" {
|
||
return nil
|
||
}
|
||
if seconds, err := strconv.ParseFloat(raw, 64); err == nil {
|
||
resetAt := now.Add(time.Duration(seconds * float64(time.Second)))
|
||
return &resetAt
|
||
}
|
||
if parsed, err := http.ParseTime(raw); err == nil {
|
||
return &parsed
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func parseOpenAIImageTryAgainCooldown(body []byte) time.Duration {
|
||
if len(body) == 0 {
|
||
return 0
|
||
}
|
||
match := openAIImageTryAgainPattern.FindSubmatch(body)
|
||
if len(match) != 3 {
|
||
return 0
|
||
}
|
||
value, err := strconv.ParseFloat(string(match[1]), 64)
|
||
if err != nil || value <= 0 {
|
||
return 0
|
||
}
|
||
switch strings.ToLower(string(match[2])) {
|
||
case "ms":
|
||
return time.Duration(value * float64(time.Millisecond))
|
||
case "s", "sec", "secs", "second", "seconds":
|
||
return time.Duration(value * float64(time.Second))
|
||
case "m", "min", "mins", "minute", "minutes":
|
||
return time.Duration(value * float64(time.Minute))
|
||
default:
|
||
return 0
|
||
}
|
||
}
|
||
|
||
const upstreamModelNotFoundCooldown = 30 * time.Minute
|
||
const upstreamModelNotFoundReason = "upstream_404_model_not_found"
|
||
const upstreamCodexPlanGatedModelCooldown = 30 * time.Minute
|
||
const upstreamCodexPlanGatedModelReason = "upstream_400_codex_plan_gated_model"
|
||
const tempUnschedBodyMaxBytes = 64 << 10
|
||
const tempUnschedMessageMaxBytes = 2048
|
||
|
||
// HandleUpstreamModelNotFound marks the requested model as temporarily
|
||
// unavailable on the account when the upstream deterministically reports it
|
||
// cannot serve that model: a 404 model-not-found, or the Codex 400 rejecting a
|
||
// plan-gated model on a ChatGPT OAuth account. Returning true tells the caller
|
||
// to fail the current attempt over to another account; the scheduler skips the
|
||
// (account, model) pair via IsSchedulableForModelWithContext until the
|
||
// cooldown expires, instead of re-selecting an account that can never serve
|
||
// the model.
|
||
func (s *RateLimitService) HandleUpstreamModelNotFound(ctx context.Context, account *Account, requestedModel string, statusCode int, responseBody []byte) bool {
|
||
if s == nil || account == nil || s.accountRepo == nil {
|
||
return false
|
||
}
|
||
if !account.ShouldHandleErrorCode(statusCode) {
|
||
return false
|
||
}
|
||
var cooldown time.Duration
|
||
var reason string
|
||
switch {
|
||
case isUpstreamModelNotFoundError(statusCode, responseBody):
|
||
cooldown, reason = upstreamModelNotFoundCooldown, upstreamModelNotFoundReason
|
||
case isOpenAIOAuthAccount(account) && isOpenAICodexPlanGatedModelError(statusCode, responseBody):
|
||
cooldown, reason = upstreamCodexPlanGatedModelCooldown, upstreamCodexPlanGatedModelReason
|
||
default:
|
||
return false
|
||
}
|
||
modelKey := modelRateLimitKeyForUpstreamModelNotFound(ctx, account, requestedModel)
|
||
if modelKey == "" {
|
||
return false
|
||
}
|
||
if shouldSkipCodexPlanGatedImageModelCooldown(ctx, reason, requestedModel, modelKey) {
|
||
return true
|
||
}
|
||
resetAt := time.Now().Add(cooldown)
|
||
if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, modelKey, resetAt, reason); err != nil {
|
||
slog.Warn("upstream_model_not_found_set_model_rate_limit_failed", "account_id", account.ID, "model", modelKey, "reason", reason, "error", err)
|
||
return true
|
||
}
|
||
slog.Info("upstream_model_not_found_model_rate_limited", "account_id", account.ID, "model", modelKey, "reason", reason, "reset_at", resetAt)
|
||
return true
|
||
}
|
||
|
||
// shouldSkipCodexPlanGatedImageModelCooldown 判断这次 Codex plan-gated 400 是否
|
||
// 属于"图片模型被文本端点拒绝"。
|
||
//
|
||
// 这类错误是确定性的端点错配,不是账号能力缺失:同一账号通过 /v1/images/* 依然
|
||
// 能出图。在这里写 per-model 冷却,会让一次用错端点的请求把整个号池对正确的生图
|
||
// 端点下线(#4828)。
|
||
//
|
||
// 但请求本身就从 /v1/images/* 入站时不适用——那种情况下被拒说明账号确实不具备
|
||
// 该模型能力,冷却是必要的刹车:没有它,每个生图请求都会完整走一遍号池,对上游
|
||
// 形成无上界的 400 放大。
|
||
//
|
||
// 请求模型与最终冷却键都要判:冷却键走的是 account.GetMappedModel,账号可能把
|
||
// 文本别名映射到 gpt-image-*,只判请求模型会漏掉这种形态。
|
||
func shouldSkipCodexPlanGatedImageModelCooldown(ctx context.Context, reason, requestedModel, modelKey string) bool {
|
||
if reason != upstreamCodexPlanGatedModelReason {
|
||
return false
|
||
}
|
||
if OpenAIImagesEndpointFromContext(ctx) {
|
||
return false
|
||
}
|
||
return IsGPTImageGenerationModel(requestedModel) || IsGPTImageGenerationModel(modelKey)
|
||
}
|
||
|
||
func modelRateLimitKeyForUpstreamModelNotFound(ctx context.Context, account *Account, requestedModel string) string {
|
||
modelKey := strings.TrimSpace(requestedModel)
|
||
if account == nil || modelKey == "" {
|
||
return modelKey
|
||
}
|
||
if account.Platform == PlatformAntigravity {
|
||
if resolved := strings.TrimSpace(resolveFinalAntigravityModelKey(ctx, account, modelKey)); resolved != "" {
|
||
return resolved
|
||
}
|
||
return modelKey
|
||
}
|
||
if mapped := strings.TrimSpace(account.GetMappedModel(modelKey)); mapped != "" {
|
||
return mapped
|
||
}
|
||
return modelKey
|
||
}
|
||
|
||
func firstRequestedModel(requestedModel []string) string {
|
||
if len(requestedModel) == 0 {
|
||
return ""
|
||
}
|
||
return strings.TrimSpace(requestedModel[0])
|
||
}
|
||
|
||
type tempUnschedulableModelContextKey struct{}
|
||
|
||
func withTempUnschedulableModel(ctx context.Context, requestedModel []string) context.Context {
|
||
model := firstRequestedModel(requestedModel)
|
||
if model == "" {
|
||
return ctx
|
||
}
|
||
if ctx == nil {
|
||
ctx = context.Background()
|
||
}
|
||
return context.WithValue(ctx, tempUnschedulableModelContextKey{}, model)
|
||
}
|
||
|
||
func tempUnschedulableModel(ctx context.Context, requestedModel []string) string {
|
||
if model := firstRequestedModel(requestedModel); model != "" {
|
||
return model
|
||
}
|
||
if ctx == nil {
|
||
return ""
|
||
}
|
||
model, _ := ctx.Value(tempUnschedulableModelContextKey{}).(string)
|
||
return strings.TrimSpace(model)
|
||
}
|
||
|
||
type tempUnschedulableRuleMatch struct {
|
||
rule TempUnschedulableRule
|
||
ruleIndex int
|
||
matchedKeyword string
|
||
}
|
||
|
||
func matchTempUnschedulableRules(account *Account, statusCode int, responseBody []byte) []tempUnschedulableRuleMatch {
|
||
if account == nil || !account.IsTempUnschedulableEnabled() || statusCode <= 0 || len(responseBody) == 0 {
|
||
return nil
|
||
}
|
||
rules := account.GetTempUnschedulableRules()
|
||
if len(rules) == 0 {
|
||
return nil
|
||
}
|
||
body := responseBody
|
||
if len(body) > tempUnschedBodyMaxBytes {
|
||
body = body[:tempUnschedBodyMaxBytes]
|
||
}
|
||
bodyLower := strings.ToLower(string(body))
|
||
matches := make([]tempUnschedulableRuleMatch, 0, 1)
|
||
for idx, rule := range rules {
|
||
if rule.ErrorCode != statusCode || len(rule.Keywords) == 0 {
|
||
continue
|
||
}
|
||
matchedKeyword := matchTempUnschedKeyword(bodyLower, rule.Keywords)
|
||
if matchedKeyword == "" {
|
||
continue
|
||
}
|
||
matches = append(matches, tempUnschedulableRuleMatch{rule: rule, ruleIndex: idx, matchedKeyword: matchedKeyword})
|
||
}
|
||
return matches
|
||
}
|
||
|
||
func (s *RateLimitService) tryTempUnschedulable(ctx context.Context, account *Account, statusCode int, responseBody []byte, requestedModel ...string) bool {
|
||
if account == nil {
|
||
return false
|
||
}
|
||
if !account.IsTempUnschedulableEnabled() {
|
||
return false
|
||
}
|
||
// 401 首次命中可临时不可调度(给 token 刷新窗口);
|
||
// 若历史上已因 401 进入过临时不可调度,则本次应升级为 error(返回 false 交由默认错误逻辑处理)。
|
||
// Antigravity 跳过:其 401 由 applyErrorPolicy 的 temp_unschedulable_rules 自行控制,无需升级逻辑。
|
||
if statusCode == http.StatusUnauthorized && account.Platform != PlatformAntigravity {
|
||
reason := account.TempUnschedulableReason
|
||
// 缓存可能没有 reason,从 DB 回退读取
|
||
if reason == "" {
|
||
if dbAcc, err := s.accountRepo.GetByID(ctx, account.ID); err == nil && dbAcc != nil {
|
||
reason = dbAcc.TempUnschedulableReason
|
||
}
|
||
}
|
||
if wasTempUnschedByStatusCode(reason, statusCode) {
|
||
slog.Info("401_escalated_to_error", "account_id", account.ID,
|
||
"reason", "previous temp-unschedulable was also 401")
|
||
return false
|
||
}
|
||
}
|
||
for _, match := range matchTempUnschedulableRules(account, statusCode, responseBody) {
|
||
if s.triggerTempUnschedulable(ctx, account, match.rule, match.ruleIndex, statusCode, match.matchedKeyword, responseBody, tempUnschedulableModel(ctx, requestedModel)) {
|
||
return true
|
||
}
|
||
}
|
||
|
||
return false
|
||
}
|
||
|
||
func wasTempUnschedByStatusCode(reason string, statusCode int) bool {
|
||
if statusCode <= 0 {
|
||
return false
|
||
}
|
||
reason = strings.TrimSpace(reason)
|
||
if reason == "" {
|
||
return false
|
||
}
|
||
|
||
var state TempUnschedState
|
||
if err := json.Unmarshal([]byte(reason), &state); err != nil {
|
||
return false
|
||
}
|
||
return state.StatusCode == statusCode
|
||
}
|
||
|
||
func matchTempUnschedKeyword(bodyLower string, keywords []string) string {
|
||
if bodyLower == "" {
|
||
return ""
|
||
}
|
||
for _, keyword := range keywords {
|
||
k := strings.TrimSpace(keyword)
|
||
if k == "" {
|
||
continue
|
||
}
|
||
if strings.Contains(bodyLower, strings.ToLower(k)) {
|
||
return k
|
||
}
|
||
}
|
||
return ""
|
||
}
|
||
|
||
func (s *RateLimitService) triggerTempUnschedulable(ctx context.Context, account *Account, rule TempUnschedulableRule, ruleIndex int, statusCode int, matchedKeyword string, responseBody []byte, requestedModel ...string) bool {
|
||
if account == nil {
|
||
return false
|
||
}
|
||
if rule.DurationMinutes <= 0 {
|
||
return false
|
||
}
|
||
|
||
now := time.Now()
|
||
until := now.Add(time.Duration(rule.DurationMinutes) * time.Minute)
|
||
|
||
state := &TempUnschedState{
|
||
UntilUnix: until.Unix(),
|
||
TriggeredAtUnix: now.Unix(),
|
||
StatusCode: statusCode,
|
||
MatchedKeyword: matchedKeyword,
|
||
RuleIndex: ruleIndex,
|
||
ErrorMessage: truncateTempUnschedMessage(responseBody, tempUnschedMessageMaxBytes),
|
||
}
|
||
|
||
reason := ""
|
||
if raw, err := json.Marshal(state); err == nil {
|
||
reason = string(raw)
|
||
}
|
||
if reason == "" {
|
||
reason = strings.TrimSpace(state.ErrorMessage)
|
||
}
|
||
|
||
// Persist known-model failures under the model key so the scheduler excludes
|
||
// only this (account, model) pair. Authentication and model-unknown failures
|
||
// retain the legacy account-wide temporary-unschedulable behavior below.
|
||
modelKey := firstRequestedModel(requestedModel)
|
||
if modelKey != "" && statusCode != http.StatusUnauthorized {
|
||
if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, modelKey, until, reason); err != nil {
|
||
slog.Warn("temp_unsched_model_rate_limit_set_failed", "account_id", account.ID, "model", modelKey, "error", err)
|
||
// The rule matched, so fail over the current request even if persistence
|
||
// failed; never widen a model-scoped failure into an account-wide block.
|
||
return true
|
||
}
|
||
slog.Info("account_model_temp_unschedulable", "account_id", account.ID, "model", modelKey, "until", until, "rule_index", ruleIndex, "status_code", statusCode)
|
||
return true
|
||
}
|
||
|
||
s.notifyAccountSchedulingBlocked(account, until, "temp_unschedulable")
|
||
if err := s.accountRepo.SetTempUnschedulable(ctx, account.ID, until, reason); err != nil {
|
||
slog.Warn("temp_unsched_set_failed", "account_id", account.ID, "error", err)
|
||
return false
|
||
}
|
||
|
||
if s.tempUnschedCache != nil {
|
||
if err := s.tempUnschedCache.SetTempUnsched(ctx, account.ID, state); err != nil {
|
||
slog.Warn("temp_unsched_cache_set_failed", "account_id", account.ID, "error", err)
|
||
}
|
||
}
|
||
|
||
slog.Info("account_temp_unschedulable", "account_id", account.ID, "until", until, "rule_index", ruleIndex, "status_code", statusCode)
|
||
return true
|
||
}
|
||
|
||
func truncateTempUnschedMessage(body []byte, maxBytes int) string {
|
||
if maxBytes <= 0 || len(body) == 0 {
|
||
return ""
|
||
}
|
||
if len(body) > maxBytes {
|
||
body = body[:maxBytes]
|
||
}
|
||
return strings.TrimSpace(string(body))
|
||
}
|
||
|
||
// HandleStreamTimeout 处理流数据超时
|
||
// 根据系统设置决定是否标记账户为临时不可调度或错误状态
|
||
// 返回是否应该停止该账号的调度
|
||
func (s *RateLimitService) HandleStreamTimeout(ctx context.Context, account *Account, model string) bool {
|
||
if account == nil {
|
||
return false
|
||
}
|
||
|
||
// 获取系统设置
|
||
if s.settingService == nil {
|
||
slog.Warn("stream_timeout_setting_service_missing", "account_id", account.ID)
|
||
return false
|
||
}
|
||
|
||
settings, err := s.settingService.GetStreamTimeoutSettings(ctx)
|
||
if err != nil {
|
||
slog.Warn("stream_timeout_get_settings_failed", "account_id", account.ID, "error", err)
|
||
return false
|
||
}
|
||
|
||
if !settings.Enabled {
|
||
return false
|
||
}
|
||
|
||
if settings.Action == StreamTimeoutActionNone {
|
||
return false
|
||
}
|
||
|
||
// 增加超时计数
|
||
var count int64 = 1
|
||
if s.timeoutCounterCache != nil {
|
||
count, err = s.timeoutCounterCache.IncrementTimeoutCount(ctx, account.ID, settings.ThresholdWindowMinutes)
|
||
if err != nil {
|
||
slog.Warn("stream_timeout_increment_count_failed", "account_id", account.ID, "error", err)
|
||
// 继续处理,使用 count=1
|
||
count = 1
|
||
}
|
||
}
|
||
|
||
slog.Info("stream_timeout_count", "account_id", account.ID, "count", count, "threshold", settings.ThresholdCount, "window_minutes", settings.ThresholdWindowMinutes, "model", model)
|
||
|
||
// 检查是否达到阈值
|
||
if count < int64(settings.ThresholdCount) {
|
||
return false
|
||
}
|
||
|
||
// 达到阈值,执行相应操作
|
||
switch settings.Action {
|
||
case StreamTimeoutActionTempUnsched:
|
||
return s.triggerStreamTimeoutTempUnsched(ctx, account, settings, model)
|
||
case StreamTimeoutActionError:
|
||
return s.triggerStreamTimeoutError(ctx, account, model)
|
||
default:
|
||
return false
|
||
}
|
||
}
|
||
|
||
// triggerStreamTimeoutTempUnsched 触发流超时临时不可调度
|
||
func (s *RateLimitService) triggerStreamTimeoutTempUnsched(ctx context.Context, account *Account, settings *StreamTimeoutSettings, model string) bool {
|
||
now := time.Now()
|
||
until := now.Add(time.Duration(settings.TempUnschedMinutes) * time.Minute)
|
||
|
||
state := &TempUnschedState{
|
||
UntilUnix: until.Unix(),
|
||
TriggeredAtUnix: now.Unix(),
|
||
StatusCode: 0, // 超时没有状态码
|
||
MatchedKeyword: "stream_timeout",
|
||
RuleIndex: -1, // 表示系统级规则
|
||
ErrorMessage: "Stream data interval timeout for model: " + model,
|
||
}
|
||
|
||
reason := ""
|
||
if raw, err := json.Marshal(state); err == nil {
|
||
reason = string(raw)
|
||
}
|
||
if reason == "" {
|
||
reason = state.ErrorMessage
|
||
}
|
||
|
||
s.notifyAccountSchedulingBlocked(account, until, "stream_timeout_temp_unschedulable")
|
||
if err := s.accountRepo.SetTempUnschedulable(ctx, account.ID, until, reason); err != nil {
|
||
slog.Warn("stream_timeout_set_temp_unsched_failed", "account_id", account.ID, "error", err)
|
||
return false
|
||
}
|
||
|
||
if s.tempUnschedCache != nil {
|
||
if err := s.tempUnschedCache.SetTempUnsched(ctx, account.ID, state); err != nil {
|
||
slog.Warn("stream_timeout_set_temp_unsched_cache_failed", "account_id", account.ID, "error", err)
|
||
}
|
||
}
|
||
|
||
// 重置超时计数
|
||
if s.timeoutCounterCache != nil {
|
||
if err := s.timeoutCounterCache.ResetTimeoutCount(ctx, account.ID); err != nil {
|
||
slog.Warn("stream_timeout_reset_count_failed", "account_id", account.ID, "error", err)
|
||
}
|
||
}
|
||
|
||
slog.Info("stream_timeout_temp_unschedulable", "account_id", account.ID, "until", until, "model", model)
|
||
return true
|
||
}
|
||
|
||
// triggerStreamTimeoutError 触发流超时错误状态
|
||
func (s *RateLimitService) triggerStreamTimeoutError(ctx context.Context, account *Account, model string) bool {
|
||
errorMsg := "Stream data interval timeout (repeated failures) for model: " + model
|
||
|
||
s.notifyAccountSchedulingBlocked(account, time.Time{}, "stream_timeout_error")
|
||
if err := s.accountRepo.SetError(ctx, account.ID, errorMsg); err != nil {
|
||
slog.Warn("stream_timeout_set_error_failed", "account_id", account.ID, "error", err)
|
||
return false
|
||
}
|
||
|
||
// 重置超时计数
|
||
if s.timeoutCounterCache != nil {
|
||
if err := s.timeoutCounterCache.ResetTimeoutCount(ctx, account.ID); err != nil {
|
||
slog.Warn("stream_timeout_reset_count_failed", "account_id", account.ID, "error", err)
|
||
}
|
||
}
|
||
|
||
slog.Warn("stream_timeout_account_error", "account_id", account.ID, "model", model)
|
||
return true
|
||
}
|