ratelimit_service.go 102 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111211221132114211521162117211821192120212121222123212421252126212721282129213021312132213321342135213621372138213921402141214221432144214521462147214821492150215121522153215421552156215721582159216021612162216321642165216621672168216921702171217221732174217521762177217821792180218121822183218421852186218721882189219021912192219321942195219621972198219922002201220222032204220522062207220822092210221122122213221422152216221722182219222022212222222322242225222622272228222922302231223222332234223522362237223822392240224122422243224422452246224722482249225022512252225322542255225622572258225922602261226222632264226522662267226822692270227122722273227422752276227722782279228022812282228322842285228622872288228922902291229222932294229522962297229822992300230123022303230423052306230723082309231023112312231323142315231623172318231923202321232223232324232523262327232823292330233123322333233423352336233723382339234023412342234323442345234623472348234923502351235223532354235523562357235823592360236123622363236423652366236723682369237023712372237323742375237623772378237923802381238223832384238523862387238823892390239123922393239423952396239723982399240024012402240324042405240624072408240924102411241224132414241524162417241824192420242124222423242424252426242724282429243024312432243324342435243624372438243924402441244224432444244524462447244824492450245124522453245424552456245724582459246024612462246324642465246624672468246924702471247224732474247524762477247824792480248124822483248424852486248724882489249024912492249324942495249624972498249925002501250225032504250525062507250825092510251125122513251425152516251725182519252025212522252325242525252625272528252925302531253225332534253525362537253825392540254125422543254425452546254725482549255025512552255325542555255625572558255925602561256225632564256525662567256825692570257125722573257425752576257725782579258025812582258325842585258625872588258925902591259225932594259525962597259825992600260126022603260426052606260726082609261026112612261326142615261626172618261926202621262226232624262526262627262826292630263126322633263426352636263726382639264026412642264326442645264626472648264926502651265226532654265526562657265826592660266126622663266426652666266726682669267026712672267326742675267626772678267926802681268226832684268526862687268826892690269126922693269426952696269726982699270027012702270327042705270627072708270927102711271227132714271527162717271827192720272127222723272427252726272727282729273027312732273327342735273627372738273927402741274227432744274527462747274827492750275127522753275427552756275727582759276027612762276327642765276627672768276927702771277227732774277527762777277827792780278127822783278427852786278727882789279027912792279327942795279627972798279928002801280228032804280528062807280828092810281128122813281428152816281728182819282028212822282328242825282628272828282928302831283228332834283528362837283828392840284128422843284428452846284728482849
  1. package service
  2. import (
  3. "bytes"
  4. "context"
  5. "encoding/json"
  6. "fmt"
  7. "log/slog"
  8. "net/http"
  9. "regexp"
  10. "strconv"
  11. "strings"
  12. "sync"
  13. "time"
  14. "github.com/Wei-Shaw/sub2api/internal/config"
  15. "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
  16. "github.com/tidwall/gjson"
  17. )
  18. // RateLimitService 处理限流和过载状态管理
  19. type RateLimitService struct {
  20. accountRepo AccountRepository
  21. usageRepo UsageLogRepository
  22. cfg *config.Config
  23. geminiQuotaService *GeminiQuotaService
  24. tempUnschedCache TempUnschedCache
  25. openAIAPIKeyHealth OpenAIAPIKeyHealthCache
  26. timeoutCounterCache TimeoutCounterCache
  27. openAI403CounterCache OpenAI403CounterCache
  28. settingService *SettingService
  29. tokenCacheInvalidator TokenCacheInvalidator
  30. runtimeBlocker AccountRuntimeBlocker
  31. usageCacheMu sync.RWMutex
  32. usageCache map[int64]*geminiUsageCacheEntry
  33. // OpenAI Team 联动熔断的进程内去重:teamID → 去重窗口截止时间
  34. openaiTeamLinkedMu sync.Mutex
  35. openaiTeamLinkedRecent map[string]time.Time
  36. }
  37. type AccountRuntimeBlocker interface {
  38. BlockAccountScheduling(account *Account, until time.Time, reason string)
  39. ClearAccountSchedulingBlock(accountID int64)
  40. }
  41. // SuccessfulTestRecoveryResult 表示测试成功后恢复了哪些运行时状态。
  42. type SuccessfulTestRecoveryResult struct {
  43. ClearedError bool
  44. ClearedRateLimit bool
  45. }
  46. // AccountRecoveryOptions 控制账号恢复时的附加行为。
  47. type AccountRecoveryOptions struct {
  48. InvalidateToken bool
  49. }
  50. type geminiUsageCacheEntry struct {
  51. windowStart time.Time
  52. cachedAt time.Time
  53. totals GeminiUsageTotals
  54. }
  55. type geminiUsageTotalsBatchProvider interface {
  56. GetGeminiUsageTotalsBatch(ctx context.Context, accountIDs []int64, startTime, endTime time.Time) (map[int64]GeminiUsageTotals, error)
  57. }
  58. const geminiPrecheckCacheTTL = time.Minute
  59. const (
  60. defaultRateLimit429CooldownSeconds = 5
  61. maxRateLimit429CooldownSeconds = 7200
  62. )
  63. const (
  64. openAIImageRateLimitDefaultCooldown = time.Minute
  65. openAIImageRateLimitReason = "openai_image_rate_limited"
  66. openAIImageCapabilityLossCooldown = 30 * time.Minute
  67. openAIImageCapabilityLossReason = "openai_image_capability_lost"
  68. )
  69. 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)`)
  70. var openCodeGoUsageLimitResetPattern = regexp.MustCompile(`(?i)\bresets\s+in\s+`)
  71. var openCodeGoUsageLimitDurationPartPattern = regexp.MustCompile(`(?i)^([0-9]+(?:\.[0-9]+)?)\s*(s|sec|secs|second|seconds|m|min|mins|minute|minutes|h|hr|hrs|hour|hours|d|day|days|w|week|weeks)\b`)
  72. const (
  73. openAI403CooldownMinutesDefault = 10
  74. openAI403DisableThreshold = 3
  75. openAI403CounterWindowMinutes = 180
  76. )
  77. // NewRateLimitService 创建RateLimitService实例
  78. func NewRateLimitService(accountRepo AccountRepository, usageRepo UsageLogRepository, cfg *config.Config, geminiQuotaService *GeminiQuotaService, tempUnschedCache TempUnschedCache) *RateLimitService {
  79. return &RateLimitService{
  80. accountRepo: accountRepo,
  81. usageRepo: usageRepo,
  82. cfg: cfg,
  83. geminiQuotaService: geminiQuotaService,
  84. tempUnschedCache: tempUnschedCache,
  85. usageCache: make(map[int64]*geminiUsageCacheEntry),
  86. }
  87. }
  88. // SetTimeoutCounterCache 设置超时计数器缓存(可选依赖)
  89. func (s *RateLimitService) SetTimeoutCounterCache(cache TimeoutCounterCache) {
  90. s.timeoutCounterCache = cache
  91. }
  92. func (s *RateLimitService) SetOpenAIAPIKeyHealthCache(cache OpenAIAPIKeyHealthCache) {
  93. s.openAIAPIKeyHealth = cache
  94. }
  95. // SetOpenAI403CounterCache 设置 OpenAI 403 连续失败计数器(可选依赖)
  96. func (s *RateLimitService) SetOpenAI403CounterCache(cache OpenAI403CounterCache) {
  97. s.openAI403CounterCache = cache
  98. }
  99. // SetSettingService 设置系统设置服务(可选依赖)
  100. func (s *RateLimitService) SetSettingService(settingService *SettingService) {
  101. s.settingService = settingService
  102. }
  103. // SetTokenCacheInvalidator 设置 token 缓存清理器(可选依赖)
  104. func (s *RateLimitService) SetTokenCacheInvalidator(invalidator TokenCacheInvalidator) {
  105. s.tokenCacheInvalidator = invalidator
  106. }
  107. func (s *RateLimitService) SetAccountRuntimeBlocker(blocker AccountRuntimeBlocker) {
  108. s.runtimeBlocker = blocker
  109. }
  110. func (s *RateLimitService) IsOpenAIAdvancedSchedulerStickyWeightedEnabled(ctx context.Context) bool {
  111. if s == nil || s.settingService == nil {
  112. return false
  113. }
  114. gateway := &OpenAIGatewayService{rateLimitService: s}
  115. return gateway.isOpenAIAdvancedSchedulerStickyWeightedEnabled(ctx)
  116. }
  117. func (s *RateLimitService) notifyAccountSchedulingBlocked(account *Account, until time.Time, reason string) {
  118. if s == nil || s.runtimeBlocker == nil || account == nil {
  119. return
  120. }
  121. s.runtimeBlocker.BlockAccountScheduling(account, until, reason)
  122. }
  123. func (s *RateLimitService) notifyAccountSchedulingBlockCleared(accountID int64) {
  124. if s == nil || s.runtimeBlocker == nil || accountID <= 0 {
  125. return
  126. }
  127. s.runtimeBlocker.ClearAccountSchedulingBlock(accountID)
  128. }
  129. // ApplyAccountSchedulingThreshold evaluates admin-configured per-platform
  130. // utilization thresholds and, when breached, parks the account as temp-
  131. // unschedulable until the winning window resets. Returns true when the account
  132. // is blocked (either newly or already paused for the same threshold reason).
  133. func (s *RateLimitService) ApplyAccountSchedulingThreshold(ctx context.Context, account *Account) bool {
  134. if s == nil || s.settingService == nil || s.accountRepo == nil || account == nil || account.ID <= 0 {
  135. return false
  136. }
  137. if !account.IsActive() || !account.Schedulable {
  138. return false
  139. }
  140. now := time.Now().UTC()
  141. thresholds := s.settingService.GetAccountSchedulingThresholds(ctx)
  142. decision := EvaluateAccountSchedulingThreshold(account, thresholds, now)
  143. if !decision.ShouldPause || decision.Until == nil || !decision.Until.After(now) {
  144. s.applyAnthropicFableSchedulingThreshold(ctx, account, thresholds, now)
  145. return false
  146. }
  147. reason := BuildDetailedAccountSchedulingThresholdReason(AccountSchedulingThresholdReasonInput{
  148. Platform: decision.Platform,
  149. Window: decision.Window,
  150. Scope: decision.Scope,
  151. ThresholdPercent: decision.ThresholdPercent,
  152. UsedPercent: decision.UsedPercent,
  153. Until: *decision.Until,
  154. Now: now,
  155. })
  156. if accountHasSameSchedulingThresholdPause(account, *decision.Until, reason) {
  157. return true
  158. }
  159. if !account.IsSchedulable() {
  160. return false
  161. }
  162. account.TempUnschedulableUntil = cloneTimePtr(decision.Until)
  163. account.TempUnschedulableReason = reason
  164. s.notifyAccountSchedulingBlocked(account, *decision.Until, "account_scheduling_threshold")
  165. if err := s.accountRepo.SetTempUnschedulable(ctx, account.ID, *decision.Until, reason); err != nil {
  166. slog.Warn("account_scheduling_threshold_set_temp_unsched_failed",
  167. "account_id", account.ID,
  168. "platform", decision.Platform,
  169. "window", decision.Window,
  170. "scope", decision.Scope,
  171. "threshold_percent", decision.ThresholdPercent,
  172. "used_percent", decision.UsedPercent,
  173. "until", decision.Until.UTC(),
  174. "error", err)
  175. } else if s.tempUnschedCache != nil {
  176. if state := tempUnschedStateFromStoredReason(reason, decision.Until.Unix()); state != nil {
  177. if err := s.tempUnschedCache.SetTempUnsched(ctx, account.ID, state); err != nil {
  178. slog.Warn("account_scheduling_threshold_cache_set_failed", "account_id", account.ID, "error", err)
  179. }
  180. }
  181. }
  182. slog.Info("account_scheduling_threshold_temp_unschedulable",
  183. "account_id", account.ID,
  184. "platform", decision.Platform,
  185. "window", decision.Window,
  186. "scope", decision.Scope,
  187. "threshold_percent", decision.ThresholdPercent,
  188. "used_percent", decision.UsedPercent,
  189. "until", decision.Until.UTC())
  190. return true
  191. }
  192. func (s *RateLimitService) applyAnthropicFableSchedulingThreshold(ctx context.Context, account *Account, thresholds map[string]int, now time.Time) {
  193. decision := evaluateAnthropicFableSchedulingThreshold(account, thresholds, now)
  194. if !decision.ShouldPause || decision.Until == nil || !decision.Until.After(now) {
  195. return
  196. }
  197. if account.isRateLimitActiveForKey(anthropicFableRateLimitKey) {
  198. return
  199. }
  200. reason := BuildDetailedAccountSchedulingThresholdReason(AccountSchedulingThresholdReasonInput{
  201. Platform: decision.Platform,
  202. Window: decision.Window,
  203. Scope: decision.Scope,
  204. ThresholdPercent: decision.ThresholdPercent,
  205. UsedPercent: decision.UsedPercent,
  206. Until: *decision.Until,
  207. Now: now,
  208. })
  209. setAccountModelRateLimitSnapshot(account, anthropicFableRateLimitKey, *decision.Until, reason, now)
  210. if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, anthropicFableRateLimitKey, *decision.Until, reason); err != nil {
  211. slog.Warn("anthropic_fable_scheduling_threshold_set_model_limit_failed",
  212. "account_id", account.ID,
  213. "threshold_percent", decision.ThresholdPercent,
  214. "used_percent", decision.UsedPercent,
  215. "until", decision.Until.UTC(),
  216. "error", err)
  217. return
  218. }
  219. slog.Info("anthropic_fable_scheduling_threshold_model_limited",
  220. "account_id", account.ID,
  221. "scope", anthropicFableRateLimitKey,
  222. "threshold_percent", decision.ThresholdPercent,
  223. "used_percent", decision.UsedPercent,
  224. "until", decision.Until.UTC())
  225. }
  226. func accountHasSameSchedulingThresholdPause(account *Account, until time.Time, reason string) bool {
  227. if account == nil || account.TempUnschedulableUntil == nil {
  228. return false
  229. }
  230. if account.TempUnschedulableUntil.UTC().Unix() != until.UTC().Unix() {
  231. return false
  232. }
  233. existing, ok := parseTempUnschedReasonPayload(account.TempUnschedulableReason)
  234. if !ok || existing.Source != AccountSchedulingThresholdReasonSource {
  235. return false
  236. }
  237. next, ok := parseTempUnschedReasonPayload(reason)
  238. if !ok || next.Source != AccountSchedulingThresholdReasonSource {
  239. return false
  240. }
  241. existing.TriggeredAtUnix = 0
  242. next.TriggeredAtUnix = 0
  243. return existing == next
  244. }
  245. // ErrorPolicyResult 表示错误策略检查的结果
  246. type ErrorPolicyResult int
  247. const (
  248. ErrorPolicyNone ErrorPolicyResult = iota // 未命中任何策略,继续默认逻辑
  249. ErrorPolicySkipped // 自定义错误码开启但未命中,跳过处理
  250. ErrorPolicyMatched // 自定义错误码命中,应停止调度
  251. ErrorPolicyTempUnscheduled // 临时不可调度规则命中
  252. )
  253. // CheckErrorPolicy 检查自定义错误码和临时不可调度规则。
  254. // 自定义错误码开启时覆盖后续所有逻辑(包括临时不可调度)。
  255. func (s *RateLimitService) CheckErrorPolicy(ctx context.Context, account *Account, statusCode int, responseBody []byte, requestedModel ...string) ErrorPolicyResult {
  256. ctx = withTempUnschedulableModel(ctx, requestedModel)
  257. if account.IsCustomErrorCodesEnabled() {
  258. if account.ShouldHandleErrorCode(statusCode) {
  259. return ErrorPolicyMatched
  260. }
  261. slog.Info("account_error_code_skipped", "account_id", account.ID, "status_code", statusCode)
  262. return ErrorPolicySkipped
  263. }
  264. if account.IsPoolMode() {
  265. // 池模式只跳过默认账号状态处理;管理员显式配置的临时不可调度规则仍应生效。
  266. // 401 保留现有认证错误语义,避免改变重复 401 的升级行为。
  267. if statusCode != http.StatusUnauthorized && s.tryTempUnschedulable(ctx, account, statusCode, responseBody) {
  268. return ErrorPolicyTempUnscheduled
  269. }
  270. return ErrorPolicySkipped
  271. }
  272. // The global overload cooldown is the default for ordinary accounts. Explicit
  273. // account policies above retain precedence over this fallback.
  274. if statusCode == 529 {
  275. return ErrorPolicyMatched
  276. }
  277. if s.tryTempUnschedulable(ctx, account, statusCode, responseBody, firstRequestedModel(requestedModel)) {
  278. return ErrorPolicyTempUnscheduled
  279. }
  280. return ErrorPolicyNone
  281. }
  282. // HandleUpstreamError 处理上游错误响应,标记账号状态
  283. // 返回是否应该停止该账号的调度
  284. func (s *RateLimitService) HandleUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte, requestedModel ...string) (shouldDisable bool) {
  285. ctx = withTempUnschedulableModel(ctx, requestedModel)
  286. // Team 联动熔断必须先于池模式/自定义错误码/临时不可调度的各类早退;
  287. // 同请求内与 fastpath 调用点的重复触发由方法内去重吸收。
  288. s.maybeHandleOpenAITeamLinkedError(ctx, account, statusCode, responseBody)
  289. customErrorCodesEnabled := account.IsCustomErrorCodesEnabled()
  290. // 池模式默认不标记本地账号状态;但管理员显式配置的临时不可调度规则优先。
  291. // 401 保留现有认证错误语义,不在这里改变池模式的认证处理。
  292. if account.IsPoolMode() && !customErrorCodesEnabled {
  293. if statusCode != http.StatusUnauthorized && s.tryTempUnschedulable(ctx, account, statusCode, responseBody) {
  294. return true
  295. }
  296. slog.Info("pool_mode_error_skipped", "account_id", account.ID, "status_code", statusCode)
  297. return false
  298. }
  299. // apikey 类型账号:检查自定义错误码配置
  300. // 如果启用且错误码不在列表中,则不处理(不停止调度、不标记限流/过载)
  301. if !account.ShouldHandleErrorCode(statusCode) {
  302. slog.Info("account_error_code_skipped", "account_id", account.ID, "status_code", statusCode)
  303. return false
  304. }
  305. if statusCode == 529 {
  306. if customErrorCodesEnabled {
  307. s.handleCustomErrorCode(ctx, account, statusCode, extractUpstreamErrorMessage(responseBody))
  308. return true
  309. }
  310. s.handle529(ctx, account)
  311. return false
  312. }
  313. if len(requestedModel) > 0 && s.HandleUpstreamModelNotFound(ctx, account, requestedModel[0], statusCode, responseBody) {
  314. return true
  315. }
  316. // Anthropic official 5h / 7d window exhaustion is a hard account limit.
  317. // It must take precedence over user-configured 429 temp-unsched rules,
  318. // otherwise a broad "rate limit" keyword rule can shorten a multi-hour
  319. // cooldown to a local temporary pause.
  320. if statusCode == http.StatusTooManyRequests && account.Platform == PlatformAnthropic {
  321. // Fable may be rejected because the organization has no usage credits for
  322. // this model. Anthropic reports that as 429, but it is a model entitlement
  323. // failure rather than a shared account window exhaustion.
  324. fableCreditsRequired := s.persistAnthropicFableCreditsRequired(ctx, account, headers, responseBody, firstRequestedModel(requestedModel))
  325. // 7d_oi 是 Fable 模型专属的 7d 窗口:只标记模型级限流,账号对其他模型仍可调度。
  326. fableLimited := s.persistAnthropicFableWindowLimit(ctx, account, headers)
  327. if s.persistAnthropicExhaustedWindowLimit(ctx, account, headers) {
  328. return false
  329. }
  330. if fableCreditsRequired || fableLimited {
  331. return false
  332. }
  333. }
  334. // 先尝试临时不可调度规则(401除外)
  335. // 如果匹配成功,直接返回,不执行后续禁用逻辑
  336. if statusCode != 401 {
  337. if s.tryTempUnschedulable(ctx, account, statusCode, responseBody, firstRequestedModel(requestedModel)) {
  338. return true
  339. }
  340. }
  341. upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(responseBody))
  342. upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
  343. if upstreamMsg != "" {
  344. upstreamMsg = truncateForLog([]byte(upstreamMsg), 512)
  345. }
  346. switch statusCode {
  347. case 400:
  348. // "organization has been disabled" → 永久禁用
  349. if strings.Contains(strings.ToLower(upstreamMsg), "organization has been disabled") {
  350. msg := "Organization disabled (400): " + upstreamMsg
  351. s.handleAuthError(ctx, account, msg)
  352. shouldDisable = true
  353. } else if account.Platform == PlatformAnthropic && strings.Contains(strings.ToLower(upstreamMsg), "credit balance") {
  354. // Anthropic API key 余额不足(语义等同 402),停止调度
  355. msg := "Credit balance exhausted (400): " + upstreamMsg
  356. s.handleAuthError(ctx, account, msg)
  357. shouldDisable = true
  358. } else if strings.Contains(strings.ToLower(upstreamMsg), "identity verification is required") {
  359. // KYC 身份验证要求 → 永久禁用,账号需完成身份验证后才能恢复
  360. msg := "Identity verification required (400): " + upstreamMsg
  361. s.handleAuthError(ctx, account, msg)
  362. shouldDisable = true
  363. }
  364. // 其他 400 错误(如参数问题)不处理,不禁用账号
  365. case 401:
  366. // 外审第9轮:Spark 影子无独立凭据,401 是母账号 token 问题——失效缓存 / refresh_token 判断 /
  367. // 永久禁用 / 临时不可调度都必须落到凭据 owner(母账号),否则影子(无 refresh_token)必中
  368. // "refresh_token missing"永久禁用分支、母账号 token cache 也不会被清,把母账号可恢复的 token
  369. // 问题变成影子永久死亡。母账号被标记 temp-unschedulable 后由 parentHealthyForShadow 级联排除影子。
  370. // 非影子时 resolveCredentialAccount 返回自身;母账号缺失/损坏(orphan 影子,罕见)时回退到原 account。
  371. authAccount := account
  372. if resolved, rerr := resolveCredentialAccount(ctx, s.accountRepo, account); rerr == nil && resolved != nil {
  373. authAccount = resolved
  374. }
  375. // OpenAI: token_invalidated / token_revoked 表示 token 被永久作废(非过期),直接标记 error
  376. openai401Code := extractUpstreamErrorCode(responseBody)
  377. if authAccount.Platform == PlatformOpenAI && (openai401Code == "token_invalidated" || openai401Code == "token_revoked") {
  378. msg := "Token revoked (401): account authentication permanently revoked"
  379. if upstreamMsg != "" {
  380. msg = "Token revoked (401): " + upstreamMsg
  381. }
  382. s.handleAuthError(ctx, authAccount, msg)
  383. shouldDisable = true
  384. break
  385. }
  386. // OpenAI: {"detail":"Unauthorized"} 表示 token 完全无效(非标准 OpenAI 错误格式),直接标记 error
  387. if authAccount.Platform == PlatformOpenAI && gjson.GetBytes(responseBody, "detail").String() == "Unauthorized" {
  388. msg := "Unauthorized (401): account authentication failed permanently"
  389. if upstreamMsg != "" {
  390. msg = "Unauthorized (401): " + upstreamMsg
  391. }
  392. s.handleAuthError(ctx, authAccount, msg)
  393. shouldDisable = true
  394. break
  395. }
  396. // OAuth 账号在 401 错误时临时不可调度(给 token 刷新窗口);非 OAuth 账号保持原有 SetError 行为。
  397. if authAccount.Type == AccountTypeOAuth {
  398. // 1. 失效缓存
  399. if s.tokenCacheInvalidator != nil {
  400. if err := s.tokenCacheInvalidator.InvalidateToken(ctx, authAccount); err != nil {
  401. slog.Warn("oauth_401_invalidate_cache_failed", "account_id", authAccount.ID, "error", err)
  402. }
  403. }
  404. // 缺少 refresh_token 的 OAuth 账号无法在冷却期内自愈(后台刷新服务也会跳过),
  405. // 直接走 SetError 永久禁用,避免冷却结束后再被选中产生一发无意义的 502。
  406. if strings.TrimSpace(authAccount.GetCredential("refresh_token")) == "" {
  407. msg := "Authentication failed (401): refresh_token missing, cannot recover"
  408. if upstreamMsg != "" {
  409. msg = "OAuth 401 (no refresh_token): " + upstreamMsg
  410. }
  411. s.handleAuthError(ctx, authAccount, msg)
  412. shouldDisable = true
  413. break
  414. }
  415. // 2. 临时不可调度,替代 SetError(保持 status=active 让刷新服务能拾取)
  416. // 注意:此处不再写回 account.Credentials/expires_at。
  417. // 原实现使用请求开始时的 account 快照整列覆盖 credentials JSONB(见
  418. // persistAccountCredentials → accountRepository.UpdateCredentials → SetCredentials),
  419. // 在另一个 worker 刚刷新完 refresh_token 的窄窗口内会把新 refresh_token 回滚为旧值,
  420. // 导致下一周期用旧 refresh_token 调上游拿到 invalid_grant 后,
  421. // tryRecoverFromRefreshRace 重读 DB 发现 currentRT == usedRT 也救不回来,账号被错误 disable。
  422. // 这里仅依赖 InvalidateToken + SetTempUnschedulable 让账号在冷却期内不被调度,
  423. // 冷却结束后由 token_provider 的 NeedsRefresh / token_refresh_service 走带分布式锁的正路刷新。
  424. msg := "Authentication failed (401): invalid or expired credentials"
  425. if upstreamMsg != "" {
  426. msg = "OAuth 401: " + upstreamMsg
  427. }
  428. if authAccount.Platform == PlatformAntigravity {
  429. extraUpdates := antigravityForceTokenRefreshExtra("401_invalid")
  430. if err := s.accountRepo.UpdateExtra(ctx, authAccount.ID, extraUpdates); err != nil {
  431. slog.Warn("antigravity_401_force_refresh_mark_failed", "account_id", authAccount.ID, "error", err)
  432. } else {
  433. if authAccount.Extra == nil {
  434. authAccount.Extra = make(map[string]any, len(extraUpdates))
  435. }
  436. for k, v := range extraUpdates {
  437. authAccount.Extra[k] = v
  438. }
  439. slog.Info("antigravity_401_force_refresh_marked", "account_id", authAccount.ID)
  440. }
  441. }
  442. cooldownMinutes := s.cfg.RateLimit.OAuth401CooldownMinutes
  443. if cooldownMinutes <= 0 {
  444. cooldownMinutes = 10
  445. }
  446. until := time.Now().Add(time.Duration(cooldownMinutes) * time.Minute)
  447. s.notifyAccountSchedulingBlocked(authAccount, until, "oauth_401")
  448. if err := s.accountRepo.SetTempUnschedulable(ctx, authAccount.ID, until, msg); err != nil {
  449. slog.Warn("oauth_401_set_temp_unschedulable_failed", "account_id", authAccount.ID, "error", err)
  450. }
  451. shouldDisable = true
  452. } else {
  453. // 非 OAuth:保持 SetError 行为
  454. msg := "Authentication failed (401): invalid or expired credentials"
  455. if upstreamMsg != "" {
  456. msg = "Authentication failed (401): " + upstreamMsg
  457. }
  458. s.handleAuthError(ctx, authAccount, msg)
  459. shouldDisable = true
  460. }
  461. case 402:
  462. // 国产供应商:余额不足是可恢复状态(充值/检测恢复后由周期任务自动解除),
  463. // 不能走 handleAuthError 永久置 status=error。改为可恢复的临时停调。
  464. if account.IsCNProvider() {
  465. s.handleCNProviderInsufficientBalance(ctx, account, upstreamMsg)
  466. shouldDisable = true
  467. break
  468. }
  469. // OpenAI: deactivated_workspace 表示工作区已停用,直接标记 error
  470. if account.Platform == PlatformOpenAI && gjson.GetBytes(responseBody, "detail.code").String() == "deactivated_workspace" {
  471. msg := "Workspace deactivated (402): workspace has been deactivated"
  472. s.handleAuthError(ctx, account, msg)
  473. shouldDisable = true
  474. break
  475. }
  476. // 支付要求:余额不足或计费问题,停止调度
  477. msg := "Payment required (402): insufficient balance or billing issue"
  478. if upstreamMsg != "" {
  479. msg = "Payment required (402): " + upstreamMsg
  480. }
  481. s.handleAuthError(ctx, account, msg)
  482. shouldDisable = true
  483. case 403:
  484. logger.LegacyPrintf(
  485. "service.ratelimit",
  486. "[HandleUpstreamErrorRaw] account_id=%d platform=%s type=%s status=403 request_id=%s cf_ray=%s upstream_msg=%s raw_body=%s",
  487. account.ID,
  488. account.Platform,
  489. account.Type,
  490. strings.TrimSpace(headers.Get("x-request-id")),
  491. strings.TrimSpace(headers.Get("cf-ray")),
  492. upstreamMsg,
  493. truncateForLog(responseBody, 1024),
  494. )
  495. shouldDisable = s.handle403(ctx, account, upstreamMsg, responseBody)
  496. case 429:
  497. s.handle429(ctx, account, headers, responseBody)
  498. shouldDisable = false
  499. case 529:
  500. // Handled after pool/custom-code policy gates above.
  501. shouldDisable = false
  502. default:
  503. // 自定义错误码启用时:在列表中的错误码都应该停止调度
  504. if customErrorCodesEnabled {
  505. msg := "Custom error code triggered"
  506. if upstreamMsg != "" {
  507. msg = upstreamMsg
  508. }
  509. s.handleCustomErrorCode(ctx, account, statusCode, msg)
  510. shouldDisable = true
  511. } else if statusCode >= 500 {
  512. // 未启用自定义错误码时:仅记录5xx错误
  513. slog.Warn("account_upstream_error", "account_id", account.ID, "status_code", statusCode)
  514. shouldDisable = false
  515. }
  516. }
  517. return shouldDisable
  518. }
  519. // PreCheckUsage proactively checks local quota before dispatching a request.
  520. // Returns false when the account should be skipped.
  521. func (s *RateLimitService) PreCheckUsage(ctx context.Context, account *Account, requestedModel string) (bool, error) {
  522. if account == nil || account.Platform != PlatformGemini {
  523. return true, nil
  524. }
  525. if s.usageRepo == nil || s.geminiQuotaService == nil {
  526. return true, nil
  527. }
  528. quota, ok := s.geminiQuotaService.QuotaForAccount(ctx, account)
  529. if !ok {
  530. return true, nil
  531. }
  532. now := time.Now()
  533. modelClass := geminiModelClassFromName(requestedModel)
  534. // 1) Daily quota precheck (RPD; resets at PST midnight)
  535. {
  536. var limit int64
  537. if quota.SharedRPD > 0 {
  538. limit = quota.SharedRPD
  539. } else {
  540. switch modelClass {
  541. case geminiModelFlash:
  542. limit = quota.FlashRPD
  543. default:
  544. limit = quota.ProRPD
  545. }
  546. }
  547. if limit > 0 {
  548. start := geminiDailyWindowStart(now)
  549. totals, ok := s.getGeminiUsageTotals(account.ID, start, now)
  550. if !ok {
  551. stats, err := s.usageRepo.GetModelStatsWithFilters(ctx, start, now, 0, 0, account.ID, 0, nil, nil, nil)
  552. if err != nil {
  553. return true, err
  554. }
  555. totals = geminiAggregateUsage(stats)
  556. s.setGeminiUsageTotals(account.ID, start, now, totals)
  557. }
  558. var used int64
  559. if quota.SharedRPD > 0 {
  560. used = totals.ProRequests + totals.FlashRequests
  561. } else {
  562. switch modelClass {
  563. case geminiModelFlash:
  564. used = totals.FlashRequests
  565. default:
  566. used = totals.ProRequests
  567. }
  568. }
  569. if used >= limit {
  570. resetAt := geminiDailyResetTime(now)
  571. // NOTE:
  572. // - This is a local precheck to reduce upstream 429s.
  573. // - Do NOT mark the account as rate-limited here; rate_limit_reset_at should reflect real upstream 429s.
  574. slog.Info("gemini_precheck_daily_quota_reached", "account_id", account.ID, "used", used, "limit", limit, "reset_at", resetAt)
  575. return false, nil
  576. }
  577. }
  578. }
  579. // 2) Minute quota precheck (RPM; fixed window current minute)
  580. {
  581. var limit int64
  582. if quota.SharedRPM > 0 {
  583. limit = quota.SharedRPM
  584. } else {
  585. switch modelClass {
  586. case geminiModelFlash:
  587. limit = quota.FlashRPM
  588. default:
  589. limit = quota.ProRPM
  590. }
  591. }
  592. if limit > 0 {
  593. start := now.Truncate(time.Minute)
  594. stats, err := s.usageRepo.GetModelStatsWithFilters(ctx, start, now, 0, 0, account.ID, 0, nil, nil, nil)
  595. if err != nil {
  596. return true, err
  597. }
  598. totals := geminiAggregateUsage(stats)
  599. var used int64
  600. if quota.SharedRPM > 0 {
  601. used = totals.ProRequests + totals.FlashRequests
  602. } else {
  603. switch modelClass {
  604. case geminiModelFlash:
  605. used = totals.FlashRequests
  606. default:
  607. used = totals.ProRequests
  608. }
  609. }
  610. if used >= limit {
  611. resetAt := start.Add(time.Minute)
  612. // Do not persist "rate limited" status from local precheck. See note above.
  613. slog.Info("gemini_precheck_minute_quota_reached", "account_id", account.ID, "used", used, "limit", limit, "reset_at", resetAt)
  614. return false, nil
  615. }
  616. }
  617. }
  618. return true, nil
  619. }
  620. // PreCheckUsageBatch performs quota precheck for multiple accounts in one request.
  621. // Returned map value=false means the account should be skipped.
  622. func (s *RateLimitService) PreCheckUsageBatch(ctx context.Context, accounts []*Account, requestedModel string) (map[int64]bool, error) {
  623. result := make(map[int64]bool, len(accounts))
  624. for _, account := range accounts {
  625. if account == nil {
  626. continue
  627. }
  628. result[account.ID] = true
  629. }
  630. if len(accounts) == 0 || requestedModel == "" {
  631. return result, nil
  632. }
  633. if s.usageRepo == nil || s.geminiQuotaService == nil {
  634. return result, nil
  635. }
  636. modelClass := geminiModelClassFromName(requestedModel)
  637. now := time.Now()
  638. dailyStart := geminiDailyWindowStart(now)
  639. minuteStart := now.Truncate(time.Minute)
  640. type quotaAccount struct {
  641. account *Account
  642. quota GeminiQuota
  643. }
  644. quotaAccounts := make([]quotaAccount, 0, len(accounts))
  645. for _, account := range accounts {
  646. if account == nil || account.Platform != PlatformGemini {
  647. continue
  648. }
  649. quota, ok := s.geminiQuotaService.QuotaForAccount(ctx, account)
  650. if !ok {
  651. continue
  652. }
  653. quotaAccounts = append(quotaAccounts, quotaAccount{
  654. account: account,
  655. quota: quota,
  656. })
  657. }
  658. if len(quotaAccounts) == 0 {
  659. return result, nil
  660. }
  661. // 1) Daily precheck (cached + batch DB fallback)
  662. dailyTotalsByID := make(map[int64]GeminiUsageTotals, len(quotaAccounts))
  663. dailyMissIDs := make([]int64, 0, len(quotaAccounts))
  664. for _, item := range quotaAccounts {
  665. limit := geminiDailyLimit(item.quota, modelClass)
  666. if limit <= 0 {
  667. continue
  668. }
  669. accountID := item.account.ID
  670. if totals, ok := s.getGeminiUsageTotals(accountID, dailyStart, now); ok {
  671. dailyTotalsByID[accountID] = totals
  672. continue
  673. }
  674. dailyMissIDs = append(dailyMissIDs, accountID)
  675. }
  676. if len(dailyMissIDs) > 0 {
  677. totalsBatch, err := s.getGeminiUsageTotalsBatch(ctx, dailyMissIDs, dailyStart, now)
  678. if err != nil {
  679. return result, err
  680. }
  681. for _, accountID := range dailyMissIDs {
  682. totals := totalsBatch[accountID]
  683. dailyTotalsByID[accountID] = totals
  684. s.setGeminiUsageTotals(accountID, dailyStart, now, totals)
  685. }
  686. }
  687. for _, item := range quotaAccounts {
  688. limit := geminiDailyLimit(item.quota, modelClass)
  689. if limit <= 0 {
  690. continue
  691. }
  692. accountID := item.account.ID
  693. used := geminiUsedRequests(item.quota, modelClass, dailyTotalsByID[accountID], true)
  694. if used >= limit {
  695. resetAt := geminiDailyResetTime(now)
  696. slog.Info("gemini_precheck_daily_quota_reached_batch", "account_id", accountID, "used", used, "limit", limit, "reset_at", resetAt)
  697. result[accountID] = false
  698. }
  699. }
  700. // 2) Minute precheck (batch DB)
  701. minuteIDs := make([]int64, 0, len(quotaAccounts))
  702. for _, item := range quotaAccounts {
  703. accountID := item.account.ID
  704. if !result[accountID] {
  705. continue
  706. }
  707. if geminiMinuteLimit(item.quota, modelClass) <= 0 {
  708. continue
  709. }
  710. minuteIDs = append(minuteIDs, accountID)
  711. }
  712. if len(minuteIDs) == 0 {
  713. return result, nil
  714. }
  715. minuteTotalsByID, err := s.getGeminiUsageTotalsBatch(ctx, minuteIDs, minuteStart, now)
  716. if err != nil {
  717. return result, err
  718. }
  719. for _, item := range quotaAccounts {
  720. accountID := item.account.ID
  721. if !result[accountID] {
  722. continue
  723. }
  724. limit := geminiMinuteLimit(item.quota, modelClass)
  725. if limit <= 0 {
  726. continue
  727. }
  728. used := geminiUsedRequests(item.quota, modelClass, minuteTotalsByID[accountID], false)
  729. if used >= limit {
  730. resetAt := minuteStart.Add(time.Minute)
  731. slog.Info("gemini_precheck_minute_quota_reached_batch", "account_id", accountID, "used", used, "limit", limit, "reset_at", resetAt)
  732. result[accountID] = false
  733. }
  734. }
  735. return result, nil
  736. }
  737. func (s *RateLimitService) getGeminiUsageTotalsBatch(ctx context.Context, accountIDs []int64, start, end time.Time) (map[int64]GeminiUsageTotals, error) {
  738. result := make(map[int64]GeminiUsageTotals, len(accountIDs))
  739. if len(accountIDs) == 0 {
  740. return result, nil
  741. }
  742. ids := make([]int64, 0, len(accountIDs))
  743. seen := make(map[int64]struct{}, len(accountIDs))
  744. for _, accountID := range accountIDs {
  745. if accountID <= 0 {
  746. continue
  747. }
  748. if _, ok := seen[accountID]; ok {
  749. continue
  750. }
  751. seen[accountID] = struct{}{}
  752. ids = append(ids, accountID)
  753. }
  754. if len(ids) == 0 {
  755. return result, nil
  756. }
  757. if batchReader, ok := s.usageRepo.(geminiUsageTotalsBatchProvider); ok {
  758. stats, err := batchReader.GetGeminiUsageTotalsBatch(ctx, ids, start, end)
  759. if err != nil {
  760. return nil, err
  761. }
  762. for _, accountID := range ids {
  763. result[accountID] = stats[accountID]
  764. }
  765. return result, nil
  766. }
  767. for _, accountID := range ids {
  768. stats, err := s.usageRepo.GetModelStatsWithFilters(ctx, start, end, 0, 0, accountID, 0, nil, nil, nil)
  769. if err != nil {
  770. return nil, err
  771. }
  772. result[accountID] = geminiAggregateUsage(stats)
  773. }
  774. return result, nil
  775. }
  776. func geminiDailyLimit(quota GeminiQuota, modelClass geminiModelClass) int64 {
  777. if quota.SharedRPD > 0 {
  778. return quota.SharedRPD
  779. }
  780. switch modelClass {
  781. case geminiModelFlash:
  782. return quota.FlashRPD
  783. default:
  784. return quota.ProRPD
  785. }
  786. }
  787. func geminiMinuteLimit(quota GeminiQuota, modelClass geminiModelClass) int64 {
  788. if quota.SharedRPM > 0 {
  789. return quota.SharedRPM
  790. }
  791. switch modelClass {
  792. case geminiModelFlash:
  793. return quota.FlashRPM
  794. default:
  795. return quota.ProRPM
  796. }
  797. }
  798. func geminiUsedRequests(quota GeminiQuota, modelClass geminiModelClass, totals GeminiUsageTotals, daily bool) int64 {
  799. if daily {
  800. if quota.SharedRPD > 0 {
  801. return totals.ProRequests + totals.FlashRequests
  802. }
  803. } else {
  804. if quota.SharedRPM > 0 {
  805. return totals.ProRequests + totals.FlashRequests
  806. }
  807. }
  808. switch modelClass {
  809. case geminiModelFlash:
  810. return totals.FlashRequests
  811. default:
  812. return totals.ProRequests
  813. }
  814. }
  815. func (s *RateLimitService) getGeminiUsageTotals(accountID int64, windowStart, now time.Time) (GeminiUsageTotals, bool) {
  816. s.usageCacheMu.RLock()
  817. defer s.usageCacheMu.RUnlock()
  818. if s.usageCache == nil {
  819. return GeminiUsageTotals{}, false
  820. }
  821. entry, ok := s.usageCache[accountID]
  822. if !ok || entry == nil {
  823. return GeminiUsageTotals{}, false
  824. }
  825. if !entry.windowStart.Equal(windowStart) {
  826. return GeminiUsageTotals{}, false
  827. }
  828. if now.Sub(entry.cachedAt) >= geminiPrecheckCacheTTL {
  829. return GeminiUsageTotals{}, false
  830. }
  831. return entry.totals, true
  832. }
  833. func (s *RateLimitService) setGeminiUsageTotals(accountID int64, windowStart, now time.Time, totals GeminiUsageTotals) {
  834. s.usageCacheMu.Lock()
  835. defer s.usageCacheMu.Unlock()
  836. if s.usageCache == nil {
  837. s.usageCache = make(map[int64]*geminiUsageCacheEntry)
  838. }
  839. s.usageCache[accountID] = &geminiUsageCacheEntry{
  840. windowStart: windowStart,
  841. cachedAt: now,
  842. totals: totals,
  843. }
  844. }
  845. // GeminiCooldown returns the fallback cooldown duration for Gemini 429s based on tier.
  846. func (s *RateLimitService) GeminiCooldown(ctx context.Context, account *Account) time.Duration {
  847. if account == nil {
  848. return 5 * time.Minute
  849. }
  850. if s.geminiQuotaService == nil {
  851. return 5 * time.Minute
  852. }
  853. return s.geminiQuotaService.CooldownForAccount(ctx, account)
  854. }
  855. // handleAuthError 处理认证类错误(401/403),停止账号调度
  856. func (s *RateLimitService) handleAuthError(ctx context.Context, account *Account, errorMsg string) {
  857. s.notifyAccountSchedulingBlocked(account, time.Time{}, "auth_error")
  858. if err := s.accountRepo.SetError(ctx, account.ID, errorMsg); err != nil {
  859. slog.Warn("account_set_error_failed", "account_id", account.ID, "error", err)
  860. return
  861. }
  862. slog.Warn("account_disabled_auth_error", "account_id", account.ID, "error", errorMsg)
  863. }
  864. func buildForbiddenErrorMessage(prefix string, upstreamMsg string, responseBody []byte, fallback string) string {
  865. prefix = strings.TrimSpace(prefix)
  866. if prefix != "" && !strings.HasSuffix(prefix, " ") {
  867. prefix += " "
  868. }
  869. if msg := strings.TrimSpace(upstreamMsg); msg != "" {
  870. return prefix + msg
  871. }
  872. rawBody := bytes.TrimSpace(responseBody)
  873. if len(rawBody) > 0 {
  874. if json.Valid(rawBody) {
  875. var compact bytes.Buffer
  876. if err := json.Compact(&compact, rawBody); err == nil {
  877. return prefix + truncateForLog(compact.Bytes(), 512)
  878. }
  879. }
  880. return prefix + truncateForLog(rawBody, 512)
  881. }
  882. return prefix + fallback
  883. }
  884. // handle403 处理 403 Forbidden 错误
  885. // Antigravity 平台区分 validation/violation/generic 三种类型,均 SetError 永久禁用;
  886. // 其他平台保持原有 SetError 行为。
  887. func (s *RateLimitService) handle403(ctx context.Context, account *Account, upstreamMsg string, responseBody []byte) (shouldDisable bool) {
  888. if account.Platform == PlatformAntigravity {
  889. return s.handleAntigravity403(ctx, account, upstreamMsg, responseBody)
  890. }
  891. // Kimi reports its transient per-account concurrency/business limit as a 403.
  892. // Keep the normal 403 failover signal (true), but never feed this exact message
  893. // into the escalating 403 counter that can permanently mark the account error.
  894. if isCNProviderConcurrencyLimit403(account, upstreamMsg) {
  895. s.handleCNProviderConcurrencyLimit403(ctx, account)
  896. return true
  897. }
  898. // 国产供应商与 openai 同口径:HTML 403(CDN/代理拦截页)不构成账号失效证据,
  899. // 且 403 在 failover 状态集里会被逐账号重放——直接 SetError 会让一个坏请求/
  900. // 一层坏代理连环永久禁用整组账号。走 HTML 豁免 + N 次累计 + 临时冷却。
  901. if account.Platform == PlatformOpenAI || IsCNProvider(account.Platform) {
  902. return s.handleOpenAI403(ctx, account, upstreamMsg, responseBody)
  903. }
  904. // 非 Antigravity 平台:保持原有行为
  905. msg := buildForbiddenErrorMessage(
  906. "Access forbidden (403):",
  907. upstreamMsg,
  908. responseBody,
  909. "account may be suspended or lack permissions",
  910. )
  911. s.handleAuthError(ctx, account, msg)
  912. return true
  913. }
  914. func (s *RateLimitService) handleOpenAI403(ctx context.Context, account *Account, upstreamMsg string, responseBody []byte) (shouldDisable bool) {
  915. // 上游代理 / CDN 在请求到达 OpenAI API 之前就拦下时,回的是 HTML 403 页面而不是
  916. // {"error":{...}} 结构化错误。这类响应描述的是「这条链路 / 这个端点被挡了」,
  917. // 不构成账号凭据或权限失效的证据——例如无效的 /v1/responses 子路径(#5334)。
  918. //
  919. // 据此写账号状态会把请求级错误放大成账号级处罚:首次即 temp-unschedulable,
  920. // 连续 openAI403DisableThreshold 次直接永久禁用;而 403 又在 failover 状态集里,
  921. // 同一个坏请求会被逐个账号重放,足以把整组账号打下线。
  922. //
  923. // 与既有口径一致:count_tokens 路径的 isOpenAIOAuthInputTokensUnsupported 已把
  924. // 「HTML 403 page without a structured error」按端点级响应处理;
  925. // shouldApplyOpenAIAlphaSearchAccountErrorSideEffects 的不变式也是端点级错误
  926. // 只换号、不写账号错误状态。这里只跳过账号处罚,不改变 failover 行为——
  927. // 换个走不同代理的账号仍有可能成功。
  928. if isHTMLResponse(responseBody) {
  929. slog.Warn(
  930. "openai_403_html_body_skips_account_penalty",
  931. "account_id", account.ID,
  932. "upstream_message", upstreamMsg,
  933. )
  934. return false
  935. }
  936. msg := buildForbiddenErrorMessage(
  937. "Access forbidden (403):",
  938. upstreamMsg,
  939. responseBody,
  940. "account may be suspended or lack permissions",
  941. )
  942. if s.openAI403CounterCache == nil {
  943. s.handleAuthError(ctx, account, msg)
  944. return true
  945. }
  946. count, err := s.openAI403CounterCache.IncrementOpenAI403Count(ctx, account.ID, openAI403CounterWindowMinutes)
  947. if err != nil {
  948. slog.Warn("openai_403_increment_failed", "account_id", account.ID, "error", err)
  949. s.handleAuthError(ctx, account, msg)
  950. return true
  951. }
  952. if count >= openAI403DisableThreshold {
  953. msg = fmt.Sprintf("%s | consecutive_403=%d/%d", msg, count, openAI403DisableThreshold)
  954. s.handleAuthError(ctx, account, msg)
  955. return true
  956. }
  957. until := time.Now().Add(time.Duration(openAI403CooldownMinutesDefault) * time.Minute)
  958. reason := fmt.Sprintf("OpenAI 403 temporary cooldown (%d/%d): %s", count, openAI403DisableThreshold, msg)
  959. s.notifyAccountSchedulingBlocked(account, until, "openai_403_temp")
  960. if err := s.accountRepo.SetTempUnschedulable(ctx, account.ID, until, reason); err != nil {
  961. slog.Warn("openai_403_set_temp_unschedulable_failed", "account_id", account.ID, "error", err)
  962. s.handleAuthError(ctx, account, msg)
  963. return true
  964. }
  965. slog.Warn(
  966. "openai_403_temp_unschedulable",
  967. "account_id", account.ID,
  968. "until", until,
  969. "count", count,
  970. "threshold", openAI403DisableThreshold,
  971. )
  972. return true
  973. }
  974. // handleAntigravity403 处理 Antigravity 平台的 403 错误
  975. // validation(需要验证)→ 永久 SetError(需人工去 Google 验证后恢复)
  976. // violation(违规封号)→ 永久 SetError(需人工处理)
  977. // generic(通用禁止)→ 永久 SetError
  978. func (s *RateLimitService) handleAntigravity403(ctx context.Context, account *Account, upstreamMsg string, responseBody []byte) (shouldDisable bool) {
  979. fbType := classifyForbiddenType(string(responseBody))
  980. switch fbType {
  981. case forbiddenTypeValidation:
  982. // VALIDATION_REQUIRED: 永久禁用,需人工去 Google 验证后手动恢复
  983. msg := buildForbiddenErrorMessage(
  984. "Validation required (403):",
  985. upstreamMsg,
  986. responseBody,
  987. "account needs Google verification",
  988. )
  989. if validationURL := extractValidationURL(string(responseBody)); validationURL != "" {
  990. msg += " | validation_url: " + validationURL
  991. }
  992. s.handleAuthError(ctx, account, msg)
  993. return true
  994. case forbiddenTypeViolation:
  995. // 违规封号: 永久禁用,需人工处理
  996. msg := buildForbiddenErrorMessage(
  997. "Account violation (403):",
  998. upstreamMsg,
  999. responseBody,
  1000. "terms of service violation",
  1001. )
  1002. s.handleAuthError(ctx, account, msg)
  1003. return true
  1004. default:
  1005. // 通用 403: 保持原有行为
  1006. msg := buildForbiddenErrorMessage(
  1007. "Access forbidden (403):",
  1008. upstreamMsg,
  1009. responseBody,
  1010. "account may be suspended or lack permissions",
  1011. )
  1012. s.handleAuthError(ctx, account, msg)
  1013. return true
  1014. }
  1015. }
  1016. // handleCustomErrorCode 处理自定义错误码,停止账号调度
  1017. func (s *RateLimitService) handleCustomErrorCode(ctx context.Context, account *Account, statusCode int, errorMsg string) {
  1018. msg := "Custom error code " + strconv.Itoa(statusCode) + ": " + errorMsg
  1019. s.notifyAccountSchedulingBlocked(account, time.Time{}, "custom_error_code")
  1020. if err := s.accountRepo.SetError(ctx, account.ID, msg); err != nil {
  1021. slog.Warn("account_set_error_failed", "account_id", account.ID, "status_code", statusCode, "error", err)
  1022. return
  1023. }
  1024. slog.Warn("account_disabled_custom_error", "account_id", account.ID, "status_code", statusCode, "error", errorMsg)
  1025. }
  1026. // handle429 处理429限流错误
  1027. // 解析响应头获取重置时间,标记账号为限流状态
  1028. func (s *RateLimitService) handle429(ctx context.Context, account *Account, headers http.Header, responseBody []byte) {
  1029. // OpenAI OAuth stays on the same account for the gateway's bounded retry
  1030. // window. Persisting a rate-limit reset on the first 429 would make the next
  1031. // retry ineligible and silently turn same-account recovery into a switch.
  1032. if account != nil && isOpenAIOAuthAccount(account) && s.runtimeBlocker != nil {
  1033. if checker, ok := s.runtimeBlocker.(interface {
  1034. ShouldRetryOpenAIOAuth429(*Account, http.Header, []byte) bool
  1035. }); ok && checker.ShouldRetryOpenAIOAuth429(account, headers, responseBody) {
  1036. return
  1037. }
  1038. }
  1039. // Spark 影子:限流/熔断状态 100% 由 QueryUsage(/wham/usage body 的 codex_bengalfox)驱动。
  1040. // /responses 的 429 携带的 x-codex-*/usage_limit_reached 是 global codex 道(plan/spec §8),
  1041. // 套到影子会把 spark 误耦合到 global 窗口——即便 spark 仍有配额也会被冷却到 global reset,
  1042. // 单影子场景直接变成无可用账号(外审第8轮 P1)。整段跳过;影子的 codex_* 仅由 account_usage 的
  1043. // QueryUsage→persistOpenAICodexProbeSnapshot 维护,枯竭由调度守卫处理。
  1044. if account.IsShadow() {
  1045. if account.ParentAccountID != nil {
  1046. notifyOpenAIAutoReset(*account.ParentAccountID)
  1047. }
  1048. return
  1049. }
  1050. // 国产供应商(kimi/zhipu/deepseek)的 429 走专用可恢复路径:余额不足 → 临时停调,
  1051. // Coding Plan 窗口耗尽 → 冷却到快照重置点。未命中则继续默认 429 逻辑。
  1052. if account.IsCNProvider() {
  1053. if s.applyCNProviderReactive429(ctx, account, headers, responseBody) {
  1054. return
  1055. }
  1056. }
  1057. // 1. OpenAI 平台:优先尝试解析 x-codex-* 响应头(用于 rate_limit_exceeded)
  1058. if account.Platform == PlatformOpenAI {
  1059. persistOpenAI429PlanType(ctx, s.accountRepo, account, responseBody)
  1060. s.persistOpenAICodexSnapshot(ctx, account, headers)
  1061. notifyOpenAIAutoReset(account.ID)
  1062. if resetAt := s.calculateOpenAI429ResetTime(headers); resetAt != nil {
  1063. s.notifyAccountSchedulingBlocked(account, *resetAt, "429")
  1064. if err := s.accountRepo.SetRateLimited(ctx, account.ID, *resetAt); err != nil {
  1065. slog.Warn("rate_limit_set_failed", "account_id", account.ID, "error", err)
  1066. return
  1067. }
  1068. slog.Info("openai_account_rate_limited", "account_id", account.ID, "reset_at", *resetAt)
  1069. return
  1070. }
  1071. }
  1072. // 2. Anthropic 平台:尝试解析 per-window 头(5h / 7d),选择实际触发的窗口
  1073. if result := calculateAnthropic429ResetTime(headers); result != nil {
  1074. s.notifyAccountSchedulingBlocked(account, result.resetAt, "429")
  1075. if err := s.accountRepo.SetRateLimited(ctx, account.ID, result.resetAt); err != nil {
  1076. slog.Warn("rate_limit_set_failed", "account_id", account.ID, "error", err)
  1077. return
  1078. }
  1079. // 更新 session window:优先使用 5h-reset 头精确计算,否则从 resetAt 反推
  1080. windowEnd := result.resetAt
  1081. if result.fiveHourReset != nil {
  1082. windowEnd = *result.fiveHourReset
  1083. }
  1084. windowStart := windowEnd.Add(-5 * time.Hour)
  1085. if err := s.accountRepo.UpdateSessionWindow(ctx, account.ID, &windowStart, &windowEnd, "rejected"); err != nil {
  1086. slog.Warn("rate_limit_update_session_window_failed", "account_id", account.ID, "error", err)
  1087. }
  1088. slog.Info("anthropic_account_rate_limited", "account_id", account.ID, "reset_at", result.resetAt, "reset_in", time.Until(result.resetAt).Truncate(time.Second))
  1089. return
  1090. }
  1091. // 3. 尝试从响应头解析重置时间(Anthropic 聚合头,向后兼容)
  1092. resetTimestamp := headers.Get("anthropic-ratelimit-unified-reset")
  1093. // 4. 如果响应头没有,尝试从响应体解析(OpenAI usage_limit_reached, Gemini)
  1094. if resetTimestamp == "" {
  1095. switch account.Platform {
  1096. case PlatformOpenAI:
  1097. // 尝试解析 OpenAI 的 usage_limit_reached 错误
  1098. if resetAt := parseOpenAIRateLimitResetTime(responseBody); resetAt != nil {
  1099. resetTime := time.Unix(*resetAt, 0)
  1100. s.notifyAccountSchedulingBlocked(account, resetTime, "429")
  1101. if err := s.accountRepo.SetRateLimited(ctx, account.ID, resetTime); err != nil {
  1102. slog.Warn("rate_limit_set_failed", "account_id", account.ID, "error", err)
  1103. return
  1104. }
  1105. slog.Info("account_rate_limited", "account_id", account.ID, "platform", account.Platform, "reset_at", resetTime, "reset_in", time.Until(resetTime).Truncate(time.Second))
  1106. return
  1107. }
  1108. case PlatformGemini, PlatformAntigravity:
  1109. // 尝试解析 Gemini 格式(用于其他平台)
  1110. if resetAt := ParseGeminiRateLimitResetTime(responseBody); resetAt != nil {
  1111. resetTime := time.Unix(*resetAt, 0)
  1112. s.notifyAccountSchedulingBlocked(account, resetTime, "429")
  1113. if err := s.accountRepo.SetRateLimited(ctx, account.ID, resetTime); err != nil {
  1114. slog.Warn("rate_limit_set_failed", "account_id", account.ID, "error", err)
  1115. return
  1116. }
  1117. slog.Info("account_rate_limited", "account_id", account.ID, "platform", account.Platform, "reset_at", resetTime, "reset_in", time.Until(resetTime).Truncate(time.Second))
  1118. return
  1119. }
  1120. }
  1121. // Anthropic 平台:没有限流重置时间的 429 可能是非真实限流(如 Extra usage required),
  1122. // 不适合按 5h/7d 窗口长时间封禁;但完全不标记会导致账号永不冷却,
  1123. // 调度器让每个请求反复撞同一批持续 429 的账号(failover 预算被白白烧掉,
  1124. // 客户端稳定收到 429)。因此同样走可配置的秒级兜底回避,管理端可调大或关闭。
  1125. if resetAt := parseRetryAfterResetTime(headers, time.Now()); resetAt != nil && resetAt.After(time.Now()) {
  1126. s.notifyAccountSchedulingBlocked(account, *resetAt, "429_retry_after")
  1127. if err := s.accountRepo.SetRateLimited(ctx, account.ID, *resetAt); err != nil {
  1128. slog.Warn("rate_limit_set_failed", "account_id", account.ID, "error", err)
  1129. }
  1130. return
  1131. }
  1132. if account.Platform == PlatformAnthropic {
  1133. slog.Warn("rate_limit_429_no_reset_time",
  1134. "account_id", account.ID,
  1135. "platform", account.Platform,
  1136. "reason", "no rate limit reset time in headers, likely not a real rate limit")
  1137. s.apply429FallbackRateLimit(ctx, account, "anthropic_no_reset_time")
  1138. return
  1139. }
  1140. // 其他平台:没有重置时间,使用可配置的秒级默认回避,避免误伤长时间不可调度。
  1141. s.apply429FallbackRateLimit(ctx, account, "no_reset_time")
  1142. return
  1143. }
  1144. // 解析Unix时间戳
  1145. ts, err := strconv.ParseInt(resetTimestamp, 10, 64)
  1146. if err != nil {
  1147. slog.Warn("rate_limit_reset_parse_failed", "reset_timestamp", resetTimestamp, "error", err)
  1148. s.apply429FallbackRateLimit(ctx, account, "reset_parse_failed")
  1149. return
  1150. }
  1151. resetAt := time.Unix(ts, 0)
  1152. // 标记限流状态
  1153. s.notifyAccountSchedulingBlocked(account, resetAt, "429")
  1154. if err := s.accountRepo.SetRateLimited(ctx, account.ID, resetAt); err != nil {
  1155. slog.Warn("rate_limit_set_failed", "account_id", account.ID, "error", err)
  1156. return
  1157. }
  1158. // 根据重置时间反推5h窗口
  1159. windowEnd := resetAt
  1160. windowStart := resetAt.Add(-5 * time.Hour)
  1161. if err := s.accountRepo.UpdateSessionWindow(ctx, account.ID, &windowStart, &windowEnd, "rejected"); err != nil {
  1162. slog.Warn("rate_limit_update_session_window_failed", "account_id", account.ID, "error", err)
  1163. }
  1164. slog.Info("account_rate_limited", "account_id", account.ID, "reset_at", resetAt)
  1165. }
  1166. func (s *RateLimitService) apply429FallbackRateLimit(ctx context.Context, account *Account, reason string) {
  1167. cooldown, enabled := s.get429FallbackCooldown(ctx, account)
  1168. if !enabled {
  1169. slog.Info("rate_limit_429_fallback_ignored", "account_id", account.ID, "platform", account.Platform, "reason", reason)
  1170. return
  1171. }
  1172. resetAt := time.Now().Add(cooldown)
  1173. slog.Warn("rate_limit_429_fallback_used", "account_id", account.ID, "platform", account.Platform, "reason", reason, "using_default", cooldown.String())
  1174. s.notifyAccountSchedulingBlocked(account, resetAt, "429_fallback")
  1175. if err := s.accountRepo.SetRateLimited(ctx, account.ID, resetAt); err != nil {
  1176. slog.Warn("rate_limit_set_failed", "account_id", account.ID, "error", err)
  1177. }
  1178. }
  1179. func (s *RateLimitService) get429FallbackCooldown(ctx context.Context, account *Account) (time.Duration, bool) {
  1180. if s.settingService != nil {
  1181. settings, err := s.settingService.GetRateLimit429CooldownSettings(ctx)
  1182. if err == nil && settings != nil {
  1183. if !settings.Enabled {
  1184. return 0, false
  1185. }
  1186. seconds := clampRateLimit429CooldownSeconds(settings.CooldownSeconds)
  1187. return time.Duration(seconds) * time.Second, true
  1188. }
  1189. slog.Warn("rate_limit_429_settings_read_failed", "account_id", account.ID, "error", err)
  1190. }
  1191. seconds := defaultRateLimit429CooldownSeconds
  1192. seconds = clampRateLimit429CooldownSeconds(seconds)
  1193. return time.Duration(seconds) * time.Second, true
  1194. }
  1195. func clampRateLimit429CooldownSeconds(seconds int) int {
  1196. if seconds < 1 {
  1197. return 1
  1198. }
  1199. if seconds > maxRateLimit429CooldownSeconds {
  1200. return maxRateLimit429CooldownSeconds
  1201. }
  1202. return seconds
  1203. }
  1204. // calculateOpenAI429ResetTime 从 OpenAI 429 响应头计算正确的重置时间
  1205. // 返回 nil 表示无法从响应头中确定重置时间
  1206. func calculateOpenAI429ResetTime(headers http.Header) *time.Time {
  1207. snapshot := ParseCodexRateLimitHeaders(headers)
  1208. if snapshot == nil {
  1209. return nil
  1210. }
  1211. normalized := snapshot.Normalize()
  1212. if normalized == nil {
  1213. return nil
  1214. }
  1215. now := time.Now()
  1216. // 判断哪个限制被触发(used_percent >= 100)
  1217. is7dExhausted := normalized.Used7dPercent != nil && *normalized.Used7dPercent >= 100
  1218. is5hExhausted := normalized.Used5hPercent != nil && *normalized.Used5hPercent >= 100
  1219. // 优先使用被触发限制的重置时间
  1220. if is7dExhausted && normalized.Reset7dSeconds != nil {
  1221. resetAt := now.Add(time.Duration(*normalized.Reset7dSeconds) * time.Second)
  1222. slog.Info("openai_429_7d_limit_exhausted", "reset_after_seconds", *normalized.Reset7dSeconds, "reset_at", resetAt)
  1223. return &resetAt
  1224. }
  1225. if is5hExhausted && normalized.Reset5hSeconds != nil {
  1226. resetAt := now.Add(time.Duration(*normalized.Reset5hSeconds) * time.Second)
  1227. slog.Info("openai_429_5h_limit_exhausted", "reset_after_seconds", *normalized.Reset5hSeconds, "reset_at", resetAt)
  1228. return &resetAt
  1229. }
  1230. // 未达到100%时,reset-after 只代表窗口信息,不能证明账号配额耗尽。
  1231. // 这类瞬时429必须回到可配置的兜底路径,避免未耗尽账号被长时间排除。
  1232. return nil
  1233. }
  1234. func (s *RateLimitService) calculateOpenAI429ResetTime(headers http.Header) *time.Time {
  1235. return calculateOpenAI429ResetTime(headers)
  1236. }
  1237. // anthropic429Result holds the parsed Anthropic 429 rate-limit information.
  1238. type anthropic429Result struct {
  1239. resetAt time.Time // The correct reset time to use for SetRateLimited
  1240. fiveHourReset *time.Time // 5h window reset timestamp (for session window calculation), nil if not available
  1241. }
  1242. type anthropicWindowLimit struct {
  1243. window string
  1244. resetAt time.Time
  1245. reason string
  1246. }
  1247. func selectAnthropicExhaustedWindow(headers http.Header, now time.Time) *anthropicWindowLimit {
  1248. reset5h, ok5hReset := parseAnthropicWindowReset(headers, "5h", now)
  1249. reset7d, ok7dReset := parseAnthropicWindowReset(headers, "7d", now)
  1250. exceeded5h := isAnthropic5hRejected(headers) || isAnthropicWindowExceeded(headers, "5h")
  1251. exceeded7d := isAnthropicWindowExceeded(headers, "7d")
  1252. if exceeded7d && ok7dReset {
  1253. return &anthropicWindowLimit{
  1254. window: "7d",
  1255. resetAt: reset7d,
  1256. reason: "anthropic_7d_window_exhausted",
  1257. }
  1258. }
  1259. if exceeded5h && ok5hReset {
  1260. return &anthropicWindowLimit{
  1261. window: "5h",
  1262. resetAt: reset5h,
  1263. reason: "anthropic_5h_window_exhausted",
  1264. }
  1265. }
  1266. return nil
  1267. }
  1268. func isAnthropic5hRejected(headers http.Header) bool {
  1269. return isAnthropicWindowRejected(headers, "5h")
  1270. }
  1271. func isAnthropicWindowRejected(headers http.Header, window string) bool {
  1272. return strings.EqualFold(strings.TrimSpace(headers.Get("anthropic-ratelimit-unified-"+window+"-status")), "rejected")
  1273. }
  1274. func parseAnthropicWindowReset(headers http.Header, window string, now time.Time) (time.Time, bool) {
  1275. maxAge := 8 * 24 * time.Hour
  1276. if window == "5h" {
  1277. maxAge = 6 * time.Hour
  1278. }
  1279. return parseAnthropicResetTimestamp(headers.Get("anthropic-ratelimit-unified-"+window+"-reset"), now, maxAge)
  1280. }
  1281. // parseAnthropicResetTimestamp 解析 Anthropic reset 头的 Unix 时间戳(自动识别毫秒),
  1282. // 并校验落在 (now, now+maxAge] 的合理区间内。
  1283. func parseAnthropicResetTimestamp(raw string, now time.Time, maxAge time.Duration) (time.Time, bool) {
  1284. raw = strings.TrimSpace(raw)
  1285. if raw == "" {
  1286. return time.Time{}, false
  1287. }
  1288. ts, err := strconv.ParseInt(raw, 10, 64)
  1289. if err != nil {
  1290. return time.Time{}, false
  1291. }
  1292. if ts > 1e11 {
  1293. ts = ts / 1000
  1294. }
  1295. resetAt := time.Unix(ts, 0)
  1296. if !resetAt.After(now) || resetAt.After(now.Add(maxAge)) {
  1297. return time.Time{}, false
  1298. }
  1299. return resetAt, true
  1300. }
  1301. func shouldPersistAnthropicWindowLimit(account *Account, limit *anthropicWindowLimit, now time.Time) bool {
  1302. if account == nil || limit == nil || !limit.resetAt.After(now) {
  1303. return false
  1304. }
  1305. if account.RateLimitResetAt == nil {
  1306. return true
  1307. }
  1308. if !account.RateLimitResetAt.After(now) {
  1309. return true
  1310. }
  1311. return limit.resetAt.After(*account.RateLimitResetAt)
  1312. }
  1313. func (s *RateLimitService) persistAnthropicExhaustedWindowLimit(ctx context.Context, account *Account, headers http.Header) bool {
  1314. if s == nil || s.accountRepo == nil || account == nil {
  1315. return false
  1316. }
  1317. now := time.Now()
  1318. limit := selectAnthropicExhaustedWindow(headers, now)
  1319. if limit == nil {
  1320. return false
  1321. }
  1322. if !shouldPersistAnthropicWindowLimit(account, limit, now) {
  1323. slog.Info("anthropic_window_rate_limit_kept",
  1324. "account_id", account.ID,
  1325. "window", limit.window,
  1326. "reset_at", limit.resetAt,
  1327. "existing_reset_at", account.RateLimitResetAt)
  1328. return true
  1329. }
  1330. s.notifyAccountSchedulingBlocked(account, limit.resetAt, limit.reason)
  1331. if err := s.accountRepo.SetRateLimited(ctx, account.ID, limit.resetAt); err != nil {
  1332. slog.Warn("anthropic_window_rate_limit_set_failed",
  1333. "account_id", account.ID,
  1334. "window", limit.window,
  1335. "reset_at", limit.resetAt,
  1336. "error", err)
  1337. return true
  1338. }
  1339. slog.Info("anthropic_window_rate_limited",
  1340. "account_id", account.ID,
  1341. "window", limit.window,
  1342. "reset_at", limit.resetAt,
  1343. "reset_in", time.Until(limit.resetAt).Truncate(time.Second))
  1344. return true
  1345. }
  1346. const (
  1347. anthropicFableWindowReason = "anthropic_7d_oi_window_exhausted"
  1348. anthropicFableCreditsRequiredReason = "anthropic_fable_credits_required"
  1349. )
  1350. // persistAnthropicFableCreditsRequired handles Anthropic's credits_required
  1351. // response for Fable. Although the upstream status is 429, this response only
  1352. // says that the organization cannot use Fable; marking the whole account rate
  1353. // limited would unnecessarily stop Sonnet, Opus, and Haiku scheduling.
  1354. func (s *RateLimitService) persistAnthropicFableCreditsRequired(ctx context.Context, account *Account, headers http.Header, responseBody []byte, requestedModel string) bool {
  1355. if s == nil || s.accountRepo == nil || account == nil {
  1356. return false
  1357. }
  1358. if !strings.EqualFold(strings.TrimSpace(gjson.GetBytes(responseBody, "error.details.error_code").String()), "credits_required") {
  1359. return false
  1360. }
  1361. model := strings.TrimSpace(gjson.GetBytes(responseBody, "error.details.model").String())
  1362. if model == "" {
  1363. model = strings.TrimSpace(requestedModel)
  1364. }
  1365. if !isAnthropicFableModel(model) {
  1366. return false
  1367. }
  1368. now := time.Now()
  1369. resetAt, ok := parseAnthropicResetTimestamp(headers.Get("anthropic-ratelimit-unified-reset"), now, 366*24*time.Hour)
  1370. if !ok {
  1371. cooldown, enabled := s.get429FallbackCooldown(ctx, account)
  1372. if !enabled {
  1373. slog.Info("anthropic_fable_credits_required_cooldown_ignored", "account_id", account.ID)
  1374. return true
  1375. }
  1376. resetAt = now.Add(cooldown)
  1377. }
  1378. if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, anthropicFableRateLimitKey, resetAt, anthropicFableCreditsRequiredReason); err != nil {
  1379. slog.Warn("anthropic_fable_credits_required_rate_limit_set_failed",
  1380. "account_id", account.ID,
  1381. "scope", anthropicFableRateLimitKey,
  1382. "reset_at", resetAt,
  1383. "error", err)
  1384. // The response is still known to be Fable-specific. Do not widen a
  1385. // persistence failure into an account-level rate limit.
  1386. return true
  1387. }
  1388. slog.Info("anthropic_fable_credits_required_model_rate_limited",
  1389. "account_id", account.ID,
  1390. "scope", anthropicFableRateLimitKey,
  1391. "reset_at", resetAt,
  1392. "reset_in", time.Until(resetAt).Truncate(time.Second))
  1393. return true
  1394. }
  1395. // selectAnthropicFableWindowLimit parses the Anthropic 7d_oi per-model window
  1396. // headers (the Fable-only 7d window, e.g. anthropic-ratelimit-unified-7d_oi-*).
  1397. // Unlike 5h/7d, exhaustion of this window only limits the Fable model family —
  1398. // the account must stay schedulable for other models.
  1399. //
  1400. // The 7d_oi surpassed-threshold header carries a float ("1.0") rather than
  1401. // "true", so exhaustion is detected via status=rejected or utilization >= 1.0.
  1402. // When the 7d_oi reset header is missing, the aggregated
  1403. // anthropic-ratelimit-unified-reset is used (it mirrors the binding claim's
  1404. // reset when 7d_oi is the representative claim).
  1405. func selectAnthropicFableWindowLimit(headers http.Header, now time.Time) *anthropicWindowLimit {
  1406. if !isAnthropicWindowRejected(headers, "7d_oi") && !isAnthropicWindowExceeded(headers, "7d_oi") {
  1407. return nil
  1408. }
  1409. resetAt, ok := parseAnthropicWindowReset(headers, "7d_oi", now)
  1410. if !ok {
  1411. resetAt, ok = parseAnthropicAggregateReset(headers, now)
  1412. }
  1413. if !ok {
  1414. return nil
  1415. }
  1416. return &anthropicWindowLimit{
  1417. window: "7d_oi",
  1418. resetAt: resetAt,
  1419. reason: anthropicFableWindowReason,
  1420. }
  1421. }
  1422. // parseAnthropicAggregateReset parses the aggregated
  1423. // anthropic-ratelimit-unified-reset header with the same sanity checks as the
  1424. // per-window variant (7d scale).
  1425. func parseAnthropicAggregateReset(headers http.Header, now time.Time) (time.Time, bool) {
  1426. return parseAnthropicResetTimestamp(headers.Get("anthropic-ratelimit-unified-reset"), now, 8*24*time.Hour)
  1427. }
  1428. // persistAnthropicFableWindowLimit marks the Fable model family as rate limited
  1429. // when the 7d_oi window is exhausted. Returns true when the 7d_oi window was the
  1430. // (or a) trigger of this 429, so the caller must not fall through to logic that
  1431. // would mark the whole account as rate limited.
  1432. func (s *RateLimitService) persistAnthropicFableWindowLimit(ctx context.Context, account *Account, headers http.Header) bool {
  1433. if s == nil || s.accountRepo == nil || account == nil {
  1434. return false
  1435. }
  1436. now := time.Now()
  1437. limit := selectAnthropicFableWindowLimit(headers, now)
  1438. if limit == nil {
  1439. return false
  1440. }
  1441. // 429 响应头本身携带最新的窗口用量(7d_oi utilization=1.0)。限流期内
  1442. // Fable 请求不再调度到该账号,若不在此处采样,7d F 进度条会冻结在
  1443. // 限流前的旧值直到窗口重置。
  1444. s.samplePassiveUsageFromHeaders(ctx, account, headers)
  1445. if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, anthropicFableRateLimitKey, limit.resetAt, limit.reason); err != nil {
  1446. slog.Warn("anthropic_fable_window_rate_limit_set_failed",
  1447. "account_id", account.ID,
  1448. "scope", anthropicFableRateLimitKey,
  1449. "reset_at", limit.resetAt,
  1450. "error", err)
  1451. return true
  1452. }
  1453. slog.Info("anthropic_fable_window_model_rate_limited",
  1454. "account_id", account.ID,
  1455. "scope", anthropicFableRateLimitKey,
  1456. "reset_at", limit.resetAt,
  1457. "reset_in", time.Until(limit.resetAt).Truncate(time.Second))
  1458. return true
  1459. }
  1460. // calculateAnthropic429ResetTime parses Anthropic's per-window rate-limit headers
  1461. // to determine which window (5h or 7d) actually triggered the 429.
  1462. //
  1463. // Headers used:
  1464. // - anthropic-ratelimit-unified-5h-utilization / anthropic-ratelimit-unified-5h-surpassed-threshold
  1465. // - anthropic-ratelimit-unified-5h-reset
  1466. // - anthropic-ratelimit-unified-7d-utilization / anthropic-ratelimit-unified-7d-surpassed-threshold
  1467. // - anthropic-ratelimit-unified-7d-reset
  1468. //
  1469. // Returns nil when the per-window headers are absent (caller should fall back to
  1470. // the aggregated anthropic-ratelimit-unified-reset header).
  1471. func calculateAnthropic429ResetTime(headers http.Header) *anthropic429Result {
  1472. reset5hStr := headers.Get("anthropic-ratelimit-unified-5h-reset")
  1473. reset7dStr := headers.Get("anthropic-ratelimit-unified-7d-reset")
  1474. if reset5hStr == "" && reset7dStr == "" {
  1475. return nil
  1476. }
  1477. var reset5h, reset7d *time.Time
  1478. if ts, err := strconv.ParseInt(reset5hStr, 10, 64); err == nil {
  1479. t := time.Unix(ts, 0)
  1480. reset5h = &t
  1481. }
  1482. if ts, err := strconv.ParseInt(reset7dStr, 10, 64); err == nil {
  1483. t := time.Unix(ts, 0)
  1484. reset7d = &t
  1485. }
  1486. is5hExceeded := isAnthropicWindowExceeded(headers, "5h")
  1487. is7dExceeded := isAnthropicWindowExceeded(headers, "7d")
  1488. slog.Info("anthropic_429_window_analysis",
  1489. "is_5h_exceeded", is5hExceeded,
  1490. "is_7d_exceeded", is7dExceeded,
  1491. "reset_5h", reset5hStr,
  1492. "reset_7d", reset7dStr,
  1493. )
  1494. // Select the correct reset time based on which window(s) are exceeded.
  1495. var chosen *time.Time
  1496. switch {
  1497. case is5hExceeded && is7dExceeded:
  1498. // Both exceeded → prefer 7d (longer cooldown), fall back to 5h
  1499. chosen = reset7d
  1500. if chosen == nil {
  1501. chosen = reset5h
  1502. }
  1503. case is5hExceeded:
  1504. chosen = reset5h
  1505. case is7dExceeded:
  1506. chosen = reset7d
  1507. default:
  1508. // Neither flag clearly exceeded — pick the sooner reset as best guess
  1509. chosen = pickSooner(reset5h, reset7d)
  1510. }
  1511. if chosen == nil {
  1512. return nil
  1513. }
  1514. return &anthropic429Result{resetAt: *chosen, fiveHourReset: reset5h}
  1515. }
  1516. // isAnthropicWindowExceeded checks whether a given Anthropic rate-limit window
  1517. // (e.g. "5h" or "7d") has been exceeded, using utilization and surpassed-threshold headers.
  1518. func isAnthropicWindowExceeded(headers http.Header, window string) bool {
  1519. prefix := "anthropic-ratelimit-unified-" + window + "-"
  1520. // Check surpassed-threshold first (most explicit signal)
  1521. if st := headers.Get(prefix + "surpassed-threshold"); strings.EqualFold(st, "true") {
  1522. return true
  1523. }
  1524. // Fall back to utilization >= 1.0
  1525. if utilStr := headers.Get(prefix + "utilization"); utilStr != "" {
  1526. if util, err := strconv.ParseFloat(utilStr, 64); err == nil && util >= 1.0-1e-9 {
  1527. // Use a small epsilon to handle floating point: treat 0.9999999... as >= 1.0
  1528. return true
  1529. }
  1530. }
  1531. return false
  1532. }
  1533. // pickSooner returns whichever of the two time pointers is earlier.
  1534. // If only one is non-nil, it is returned. If both are nil, returns nil.
  1535. func pickSooner(a, b *time.Time) *time.Time {
  1536. switch {
  1537. case a != nil && b != nil:
  1538. if a.Before(*b) {
  1539. return a
  1540. }
  1541. return b
  1542. case a != nil:
  1543. return a
  1544. default:
  1545. return b
  1546. }
  1547. }
  1548. func (s *RateLimitService) persistOpenAICodexSnapshot(ctx context.Context, account *Account, headers http.Header) {
  1549. if s == nil || s.accountRepo == nil || account == nil || headers == nil {
  1550. return
  1551. }
  1552. // spark 影子的 codex_* 仅由 QueryUsage(/wham/usage bengalfox 道)更新,不能被 /responses 的
  1553. // x-codex-* 全局头快照污染(外审第7轮 P1,与 updateCodexUsageSnapshot 同口径)。
  1554. if account.IsShadow() {
  1555. return
  1556. }
  1557. snapshot := ParseCodexRateLimitHeaders(headers)
  1558. if snapshot == nil {
  1559. return
  1560. }
  1561. updates := buildCodexUsageExtraUpdates(snapshot, time.Now())
  1562. if len(updates) == 0 {
  1563. return
  1564. }
  1565. if err := s.accountRepo.UpdateExtra(ctx, account.ID, updates); err != nil {
  1566. slog.Warn("openai_codex_snapshot_persist_failed", "account_id", account.ID, "error", err)
  1567. return
  1568. }
  1569. notifyOpenAIAutoReset(account.ID)
  1570. }
  1571. // parseOpenAIRateLimitResetTime 解析 OpenAI 兼容格式的 429 响应,返回重置时间的 Unix 时间戳
  1572. // OpenAI 的 usage_limit_reached 错误格式:
  1573. //
  1574. // {
  1575. // "error": {
  1576. // "message": "The usage limit has been reached",
  1577. // "type": "usage_limit_reached",
  1578. // "resets_at": 1769404154,
  1579. // "resets_in_seconds": 133107
  1580. // }
  1581. // }
  1582. func parseOpenAIRateLimitResetTime(body []byte) *int64 {
  1583. var parsed map[string]any
  1584. if err := json.Unmarshal(body, &parsed); err != nil {
  1585. return nil
  1586. }
  1587. errObj, ok := parsed["error"].(map[string]any)
  1588. if !ok {
  1589. return nil
  1590. }
  1591. // 检查是否为已知的账号用量限制类型。
  1592. errType, _ := errObj["type"].(string)
  1593. if errType != "usage_limit_reached" && errType != "rate_limit_exceeded" && errType != "GoUsageLimitError" {
  1594. return nil
  1595. }
  1596. // 优先使用 resets_at(Unix 时间戳)
  1597. if resetsAt, ok := errObj["resets_at"].(float64); ok {
  1598. ts := int64(resetsAt)
  1599. return &ts
  1600. }
  1601. if resetsAt, ok := errObj["resets_at"].(string); ok {
  1602. if ts, err := strconv.ParseInt(resetsAt, 10, 64); err == nil {
  1603. return &ts
  1604. }
  1605. }
  1606. // 如果没有 resets_at,尝试使用 resets_in_seconds
  1607. if resetsInSeconds, ok := errObj["resets_in_seconds"].(float64); ok {
  1608. ts := time.Now().Unix() + int64(resetsInSeconds)
  1609. return &ts
  1610. }
  1611. if resetsInSeconds, ok := errObj["resets_in_seconds"].(string); ok {
  1612. if sec, err := strconv.ParseInt(resetsInSeconds, 10, 64); err == nil {
  1613. ts := time.Now().Unix() + sec
  1614. return &ts
  1615. }
  1616. }
  1617. // OpenCode Go subscriptions expose the reset only in a human-readable message,
  1618. // for example: "Weekly usage limit reached. Resets in 2 days."
  1619. if errType == "GoUsageLimitError" {
  1620. message, _ := errObj["message"].(string)
  1621. if resetAfter := parseOpenCodeGoUsageLimitResetDuration(message); resetAfter > 0 {
  1622. ts := time.Now().Add(resetAfter).Unix()
  1623. return &ts
  1624. }
  1625. }
  1626. return nil
  1627. }
  1628. func parseOpenCodeGoUsageLimitResetDuration(message string) time.Duration {
  1629. resetPrefix := openCodeGoUsageLimitResetPattern.FindStringIndex(message)
  1630. if resetPrefix == nil {
  1631. return 0
  1632. }
  1633. remainder := message[resetPrefix[1]:]
  1634. var total time.Duration
  1635. for {
  1636. remainder = strings.TrimSpace(remainder)
  1637. matches := openCodeGoUsageLimitDurationPartPattern.FindStringSubmatchIndex(remainder)
  1638. if matches == nil {
  1639. break
  1640. }
  1641. value, err := strconv.ParseFloat(remainder[matches[2]:matches[3]], 64)
  1642. if err != nil || value <= 0 {
  1643. return 0
  1644. }
  1645. unit := openCodeGoUsageLimitDurationUnit(remainder[matches[4]:matches[5]])
  1646. if unit <= 0 {
  1647. return 0
  1648. }
  1649. const maxDuration = time.Duration(1<<63 - 1)
  1650. if value >= float64(maxDuration)/float64(unit) {
  1651. return 0
  1652. }
  1653. part := time.Duration(value * float64(unit))
  1654. if part <= 0 || total > maxDuration-part {
  1655. return 0
  1656. }
  1657. total += part
  1658. remainder = remainder[matches[1]:]
  1659. }
  1660. return total
  1661. }
  1662. func openCodeGoUsageLimitDurationUnit(raw string) time.Duration {
  1663. switch strings.ToLower(raw) {
  1664. case "s", "sec", "secs", "second", "seconds":
  1665. return time.Second
  1666. case "m", "min", "mins", "minute", "minutes":
  1667. return time.Minute
  1668. case "h", "hr", "hrs", "hour", "hours":
  1669. return time.Hour
  1670. case "d", "day", "days":
  1671. return 24 * time.Hour
  1672. case "w", "week", "weeks":
  1673. return 7 * 24 * time.Hour
  1674. default:
  1675. return 0
  1676. }
  1677. }
  1678. func parseOpenAIRateLimitPlanType(body []byte) string {
  1679. var parsed map[string]any
  1680. if err := json.Unmarshal(body, &parsed); err != nil {
  1681. return ""
  1682. }
  1683. errObj, ok := parsed["error"].(map[string]any)
  1684. if !ok {
  1685. return ""
  1686. }
  1687. errType, _ := errObj["type"].(string)
  1688. if errType != "usage_limit_reached" && errType != "rate_limit_exceeded" {
  1689. return ""
  1690. }
  1691. planType, _ := errObj["plan_type"].(string)
  1692. return strings.ToLower(strings.TrimSpace(planType))
  1693. }
  1694. func persistOpenAI429PlanType(ctx context.Context, repo AccountRepository, account *Account, body []byte) {
  1695. if repo == nil || account == nil || account.Platform != PlatformOpenAI {
  1696. return
  1697. }
  1698. // spark 影子账号恒不持凭据:即便收到带 plan_type 的 429,也不能把 plan_type 写进影子 credentials
  1699. // ——该路径走 repo.BulkUpdate 直写、不经 persistAccountCredentials 守卫(外审第7轮 P1)。
  1700. // plan_type 由母账号在自己的请求上维护,影子跳过。
  1701. if account.IsCredentialShadow() {
  1702. return
  1703. }
  1704. planType := parseOpenAIRateLimitPlanType(body)
  1705. if planType == "" {
  1706. return
  1707. }
  1708. current := strings.TrimSpace(account.GetCredential("plan_type"))
  1709. if strings.EqualFold(current, planType) {
  1710. return
  1711. }
  1712. if _, err := repo.BulkUpdate(ctx, []int64{account.ID}, AccountBulkUpdate{
  1713. Credentials: map[string]any{"plan_type": planType},
  1714. }); err != nil {
  1715. slog.Warn("openai_429_plan_type_sync_failed", "account_id", account.ID, "plan_type", planType, "error", err)
  1716. return
  1717. }
  1718. if account.Credentials == nil {
  1719. account.Credentials = make(map[string]any, 1)
  1720. }
  1721. account.Credentials["plan_type"] = planType
  1722. slog.Info("openai_429_plan_type_synced", "account_id", account.ID, "previous_plan_type", current, "plan_type", planType)
  1723. }
  1724. // handle529 处理529过载错误
  1725. // 根据配置决定是否暂停账号调度及冷却时长
  1726. func (s *RateLimitService) handle529(ctx context.Context, account *Account) {
  1727. var settings *OverloadCooldownSettings
  1728. if s.settingService != nil {
  1729. var err error
  1730. settings, err = s.settingService.GetOverloadCooldownSettings(ctx)
  1731. if err != nil {
  1732. slog.Warn("overload_settings_read_failed", "account_id", account.ID, "error", err)
  1733. settings = nil
  1734. }
  1735. }
  1736. // 回退到配置文件
  1737. if settings == nil {
  1738. cooldown := s.cfg.RateLimit.OverloadCooldownMinutes
  1739. if cooldown <= 0 {
  1740. cooldown = 10
  1741. }
  1742. settings = &OverloadCooldownSettings{Enabled: true, CooldownMinutes: cooldown}
  1743. }
  1744. if !settings.Enabled {
  1745. slog.Info("account_529_ignored", "account_id", account.ID, "reason", "overload_cooldown_disabled")
  1746. return
  1747. }
  1748. cooldownMinutes := settings.CooldownMinutes
  1749. if cooldownMinutes <= 0 {
  1750. cooldownMinutes = 10
  1751. }
  1752. until := time.Now().Add(time.Duration(cooldownMinutes) * time.Minute)
  1753. s.notifyAccountSchedulingBlocked(account, until, "529")
  1754. if err := s.accountRepo.SetOverloaded(ctx, account.ID, until); err != nil {
  1755. slog.Warn("overload_set_failed", "account_id", account.ID, "error", err)
  1756. return
  1757. }
  1758. slog.Info("account_overloaded", "account_id", account.ID, "until", until)
  1759. }
  1760. // UpdateSessionWindow 从成功响应更新5h窗口状态
  1761. func (s *RateLimitService) UpdateSessionWindow(ctx context.Context, account *Account, headers http.Header) {
  1762. status := headers.Get("anthropic-ratelimit-unified-5h-status")
  1763. if status == "" {
  1764. return
  1765. }
  1766. // 检查是否需要初始化时间窗口
  1767. // 对于 Setup Token 账号,首次成功请求时需要预测时间窗口
  1768. var windowStart, windowEnd *time.Time
  1769. needInitWindow := account.SessionWindowEnd == nil || time.Now().After(*account.SessionWindowEnd)
  1770. // 优先使用响应头中的真实重置时间(比预测更准确)
  1771. if resetStr := headers.Get("anthropic-ratelimit-unified-5h-reset"); resetStr != "" {
  1772. if ts, err := strconv.ParseInt(resetStr, 10, 64); err == nil {
  1773. // 检测可能的毫秒时间戳(秒级约为 1e9,毫秒约为 1e12)
  1774. if ts > 1e11 {
  1775. slog.Warn("account_session_window_header_millis_detected", "account_id", account.ID, "raw_reset", resetStr)
  1776. ts = ts / 1000
  1777. }
  1778. end := time.Unix(ts, 0)
  1779. // 校验时间戳是否在合理范围内(不早于 5h 前,不晚于 7 天后)
  1780. minAllowed := time.Now().Add(-5 * time.Hour)
  1781. maxAllowed := time.Now().Add(7 * 24 * time.Hour)
  1782. if end.Before(minAllowed) || end.After(maxAllowed) {
  1783. slog.Warn("account_session_window_header_out_of_range", "account_id", account.ID, "raw_reset", resetStr, "parsed_end", end)
  1784. } else if needInitWindow || account.SessionWindowEnd == nil || !end.Equal(*account.SessionWindowEnd) {
  1785. // 窗口需要初始化,或者真实重置时间与已存储的不同,则更新
  1786. start := end.Add(-5 * time.Hour)
  1787. windowStart = &start
  1788. windowEnd = &end
  1789. slog.Info("account_session_window_from_header", "account_id", account.ID, "window_start", start, "window_end", end, "status", status)
  1790. }
  1791. } else {
  1792. slog.Warn("account_session_window_header_parse_failed", "account_id", account.ID, "raw_reset", resetStr, "error", err)
  1793. }
  1794. }
  1795. // 回退:如果没有真实重置时间且需要初始化窗口,使用预测
  1796. if windowEnd == nil && needInitWindow && (status == "allowed" || status == "allowed_warning") {
  1797. now := time.Now()
  1798. start := time.Date(now.Year(), now.Month(), now.Day(), now.Hour(), 0, 0, 0, now.Location())
  1799. end := start.Add(5 * time.Hour)
  1800. windowStart = &start
  1801. windowEnd = &end
  1802. slog.Info("account_session_window_initialized", "account_id", account.ID, "window_start", start, "window_end", end, "status", status)
  1803. }
  1804. // 窗口重置时清除旧的 utilization 和被动采样数据,避免残留上个窗口的数据
  1805. if windowEnd != nil && needInitWindow {
  1806. _ = s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{
  1807. "session_window_utilization": nil,
  1808. "passive_usage_7d_utilization": nil,
  1809. "passive_usage_7d_reset": nil,
  1810. "passive_usage_7d_oi_utilization": nil,
  1811. "passive_usage_7d_oi_reset": nil,
  1812. "passive_usage_sampled_at": nil,
  1813. })
  1814. }
  1815. if err := s.accountRepo.UpdateSessionWindow(ctx, account.ID, windowStart, windowEnd, status); err != nil {
  1816. slog.Warn("session_window_update_failed", "account_id", account.ID, "error", err)
  1817. }
  1818. // 被动采样:从响应头收集 5h + 7d + 7d_oi utilization,合并为一次 DB 写入
  1819. s.samplePassiveUsageFromHeaders(ctx, account, headers)
  1820. // 如果状态为allowed且之前有限流,说明窗口已重置,清除限流状态
  1821. if status == "allowed" && account.IsRateLimited() {
  1822. if err := s.ClearRateLimit(ctx, account.ID); err != nil {
  1823. slog.Warn("rate_limit_clear_failed", "account_id", account.ID, "error", err)
  1824. }
  1825. }
  1826. }
  1827. // samplePassiveUsageFromHeaders 从 Anthropic 响应头收集 5h/7d/7d_oi 的
  1828. // utilization 与 reset 被动采样数据,合并为一次 Extra 写入。无数据时不写。
  1829. func (s *RateLimitService) samplePassiveUsageFromHeaders(ctx context.Context, account *Account, headers http.Header) {
  1830. extraUpdates := make(map[string]any, 6)
  1831. // 5h utilization(0-1 小数),供 estimateSetupTokenUsage 使用
  1832. if utilStr := headers.Get("anthropic-ratelimit-unified-5h-utilization"); utilStr != "" {
  1833. if util, err := strconv.ParseFloat(utilStr, 64); err == nil {
  1834. extraUpdates["session_window_utilization"] = util
  1835. }
  1836. }
  1837. // 7d utilization(0-1 小数)
  1838. if utilStr := headers.Get("anthropic-ratelimit-unified-7d-utilization"); utilStr != "" {
  1839. if util, err := strconv.ParseFloat(utilStr, 64); err == nil {
  1840. extraUpdates["passive_usage_7d_utilization"] = util
  1841. }
  1842. }
  1843. // 7d reset timestamp
  1844. if resetStr := headers.Get("anthropic-ratelimit-unified-7d-reset"); resetStr != "" {
  1845. if ts, err := strconv.ParseInt(resetStr, 10, 64); err == nil {
  1846. if ts > 1e11 {
  1847. ts = ts / 1000
  1848. }
  1849. extraUpdates["passive_usage_7d_reset"] = ts
  1850. }
  1851. }
  1852. // 7d_oi (Fable 专属 7d 窗口) utilization(0-1 小数)
  1853. if utilStr := headers.Get("anthropic-ratelimit-unified-7d_oi-utilization"); utilStr != "" {
  1854. if util, err := strconv.ParseFloat(utilStr, 64); err == nil {
  1855. extraUpdates["passive_usage_7d_oi_utilization"] = util
  1856. }
  1857. }
  1858. // 7d_oi reset timestamp
  1859. if resetStr := headers.Get("anthropic-ratelimit-unified-7d_oi-reset"); resetStr != "" {
  1860. if ts, err := strconv.ParseInt(resetStr, 10, 64); err == nil {
  1861. if ts > 1e11 {
  1862. ts = ts / 1000
  1863. }
  1864. extraUpdates["passive_usage_7d_oi_reset"] = ts
  1865. }
  1866. }
  1867. if len(extraUpdates) > 0 {
  1868. extraUpdates["passive_usage_sampled_at"] = time.Now().UTC().Format(time.RFC3339)
  1869. if err := s.accountRepo.UpdateExtra(ctx, account.ID, extraUpdates); err != nil {
  1870. slog.Warn("passive_usage_update_failed", "account_id", account.ID, "error", err)
  1871. }
  1872. }
  1873. }
  1874. // ClearRateLimit 清除账号的限流状态
  1875. func (s *RateLimitService) ClearRateLimit(ctx context.Context, accountID int64) error {
  1876. if err := s.accountRepo.ClearRateLimit(ctx, accountID); err != nil {
  1877. return err
  1878. }
  1879. if err := s.accountRepo.ClearAntigravityQuotaScopes(ctx, accountID); err != nil {
  1880. return err
  1881. }
  1882. if err := s.accountRepo.ClearModelRateLimits(ctx, accountID); err != nil {
  1883. return err
  1884. }
  1885. // 清除限流时一并清理临时不可调度状态,避免周限/窗口重置后仍被本地临时状态阻断。
  1886. if err := s.accountRepo.ClearTempUnschedulable(ctx, accountID); err != nil {
  1887. return err
  1888. }
  1889. if s.tempUnschedCache != nil {
  1890. if err := s.tempUnschedCache.DeleteTempUnsched(ctx, accountID); err != nil {
  1891. slog.Warn("temp_unsched_cache_delete_failed", "account_id", accountID, "error", err)
  1892. }
  1893. }
  1894. s.ResetOpenAI403Counter(ctx, accountID)
  1895. s.notifyAccountSchedulingBlockCleared(accountID)
  1896. return nil
  1897. }
  1898. func (s *RateLimitService) ResetOpenAI403Counter(ctx context.Context, accountID int64) {
  1899. if s == nil || s.openAI403CounterCache == nil || accountID <= 0 {
  1900. return
  1901. }
  1902. if err := s.openAI403CounterCache.ResetOpenAI403Count(ctx, accountID); err != nil {
  1903. slog.Warn("openai_403_reset_failed", "account_id", accountID, "error", err)
  1904. }
  1905. }
  1906. // RecoverAccountState 按需恢复账号的可恢复运行时状态。
  1907. func (s *RateLimitService) RecoverAccountState(ctx context.Context, accountID int64, options AccountRecoveryOptions) (*SuccessfulTestRecoveryResult, error) {
  1908. account, err := s.accountRepo.GetByID(ctx, accountID)
  1909. if err != nil {
  1910. return nil, err
  1911. }
  1912. result := &SuccessfulTestRecoveryResult{}
  1913. if account.Status == StatusError {
  1914. if err := s.accountRepo.ClearError(ctx, accountID); err != nil {
  1915. return nil, err
  1916. }
  1917. result.ClearedError = true
  1918. if options.InvalidateToken && s.tokenCacheInvalidator != nil && account.IsOAuth() {
  1919. if invalidateErr := s.tokenCacheInvalidator.InvalidateToken(ctx, account); invalidateErr != nil {
  1920. slog.Warn("recover_account_state_invalidate_token_failed", "account_id", accountID, "error", invalidateErr)
  1921. }
  1922. }
  1923. }
  1924. if hasRecoverableRuntimeState(account) {
  1925. if err := s.ClearRateLimit(ctx, accountID); err != nil {
  1926. return nil, err
  1927. }
  1928. result.ClearedRateLimit = true
  1929. }
  1930. if result.ClearedError || result.ClearedRateLimit {
  1931. s.ResetOpenAI403Counter(ctx, accountID)
  1932. if result.ClearedError && !result.ClearedRateLimit {
  1933. s.notifyAccountSchedulingBlockCleared(accountID)
  1934. }
  1935. }
  1936. return result, nil
  1937. }
  1938. // A successful probe describes only that request. It cannot clear unrelated
  1939. // model limits or newer failures written by another request. Administrators
  1940. // can explicitly recover state through RecoverAccountState after reviewing it.
  1941. func (s *RateLimitService) RecoverAccountAfterSuccessfulTest(ctx context.Context, accountID int64) (*SuccessfulTestRecoveryResult, error) {
  1942. if _, err := s.accountRepo.GetByID(ctx, accountID); err != nil {
  1943. return nil, err
  1944. }
  1945. return &SuccessfulTestRecoveryResult{}, nil
  1946. }
  1947. func (s *RateLimitService) ClearTempUnschedulable(ctx context.Context, accountID int64) error {
  1948. if err := s.accountRepo.ClearTempUnschedulable(ctx, accountID); err != nil {
  1949. return err
  1950. }
  1951. if s.tempUnschedCache != nil {
  1952. if err := s.tempUnschedCache.DeleteTempUnsched(ctx, accountID); err != nil {
  1953. slog.Warn("temp_unsched_cache_delete_failed", "account_id", accountID, "error", err)
  1954. }
  1955. }
  1956. // 同时清除模型级别限流
  1957. if err := s.accountRepo.ClearModelRateLimits(ctx, accountID); err != nil {
  1958. slog.Warn("clear_model_rate_limits_on_temp_unsched_reset_failed", "account_id", accountID, "error", err)
  1959. }
  1960. s.notifyAccountSchedulingBlockCleared(accountID)
  1961. return nil
  1962. }
  1963. func hasRecoverableRuntimeState(account *Account) bool {
  1964. if account == nil {
  1965. return false
  1966. }
  1967. if account.RateLimitedAt != nil || account.RateLimitResetAt != nil || account.OverloadUntil != nil || account.TempUnschedulableUntil != nil {
  1968. return true
  1969. }
  1970. if len(account.Extra) == 0 {
  1971. return false
  1972. }
  1973. return hasNonEmptyMapValue(account.Extra, "model_rate_limits") ||
  1974. hasNonEmptyMapValue(account.Extra, "antigravity_quota_scopes")
  1975. }
  1976. func hasNonEmptyMapValue(extra map[string]any, key string) bool {
  1977. raw, ok := extra[key]
  1978. if !ok || raw == nil {
  1979. return false
  1980. }
  1981. switch typed := raw.(type) {
  1982. case map[string]any:
  1983. return len(typed) > 0
  1984. case map[string]string:
  1985. return len(typed) > 0
  1986. case []any:
  1987. return len(typed) > 0
  1988. default:
  1989. return true
  1990. }
  1991. }
  1992. func (s *RateLimitService) GetTempUnschedStatus(ctx context.Context, accountID int64) (*TempUnschedState, error) {
  1993. now := time.Now().Unix()
  1994. if s.tempUnschedCache != nil {
  1995. state, err := s.tempUnschedCache.GetTempUnsched(ctx, accountID)
  1996. if err != nil {
  1997. return nil, err
  1998. }
  1999. if state != nil && state.UntilUnix > now {
  2000. return state, nil
  2001. }
  2002. }
  2003. account, err := s.accountRepo.GetByID(ctx, accountID)
  2004. if err != nil {
  2005. return nil, err
  2006. }
  2007. if account.TempUnschedulableUntil == nil {
  2008. return nil, nil
  2009. }
  2010. if account.TempUnschedulableUntil.Unix() <= now {
  2011. return nil, nil
  2012. }
  2013. state := &TempUnschedState{
  2014. UntilUnix: account.TempUnschedulableUntil.Unix(),
  2015. }
  2016. if account.TempUnschedulableReason != "" {
  2017. var parsed TempUnschedState
  2018. if err := json.Unmarshal([]byte(account.TempUnschedulableReason), &parsed); err == nil {
  2019. if parsed.UntilUnix == 0 {
  2020. parsed.UntilUnix = state.UntilUnix
  2021. }
  2022. state = &parsed
  2023. } else {
  2024. state.ErrorMessage = account.TempUnschedulableReason
  2025. }
  2026. }
  2027. if s.tempUnschedCache != nil {
  2028. if err := s.tempUnschedCache.SetTempUnsched(ctx, accountID, state); err != nil {
  2029. slog.Warn("temp_unsched_cache_set_failed", "account_id", accountID, "error", err)
  2030. }
  2031. }
  2032. return state, nil
  2033. }
  2034. func (s *RateLimitService) HandleTempUnschedulable(ctx context.Context, account *Account, statusCode int, responseBody []byte, requestedModel ...string) bool {
  2035. if account == nil {
  2036. return false
  2037. }
  2038. if account.IsPoolMode() && !account.IsCustomErrorCodesEnabled() {
  2039. return false
  2040. }
  2041. if !account.ShouldHandleErrorCode(statusCode) {
  2042. return false
  2043. }
  2044. ctx = withTempUnschedulableModel(ctx, requestedModel)
  2045. return s.tryTempUnschedulable(ctx, account, statusCode, responseBody, firstRequestedModel(requestedModel))
  2046. }
  2047. func (s *RateLimitService) HandleOpenAIImageRateLimit(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte) bool {
  2048. if s == nil || account == nil || s.accountRepo == nil {
  2049. return false
  2050. }
  2051. if account.Platform != PlatformOpenAI {
  2052. return false
  2053. }
  2054. if !account.ShouldHandleErrorCode(statusCode) {
  2055. slog.Info("openai_image_rate_limit_skipped_by_error_code_policy", "account_id", account.ID, "status_code", statusCode)
  2056. return false
  2057. }
  2058. if !isOpenAIImageRateLimitError(statusCode, responseBody) {
  2059. return false
  2060. }
  2061. resetAt := openAIImageRateLimitResetAt(headers, responseBody)
  2062. if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, openAIImageGenerationRateLimitKey, resetAt, openAIImageRateLimitReason); err != nil {
  2063. slog.Warn("openai_image_rate_limit_set_model_rate_limit_failed", "account_id", account.ID, "scope", openAIImageGenerationRateLimitKey, "error", err)
  2064. return true
  2065. }
  2066. slog.Info("openai_image_rate_limited", "account_id", account.ID, "scope", openAIImageGenerationRateLimitKey, "reset_at", resetAt, "reset_in", time.Until(resetAt).Truncate(time.Second))
  2067. return true
  2068. }
  2069. // HandleOpenAICodexSparkRateLimit 将 Spark 独立配额窗口记录为模型级限流。
  2070. // Spark 的 x-codex-* 使用率和 reset 时间只代表 Spark 模型维度,不能写入账号级
  2071. // RateLimitResetAt,否则同一 OAuth 账号上的其他模型也会被错误停调。
  2072. func (s *RateLimitService) HandleOpenAICodexSparkRateLimit(ctx context.Context, account *Account, requestedModel string, statusCode int, headers http.Header, responseBody []byte) bool {
  2073. if s == nil || account == nil || s.accountRepo == nil || statusCode != http.StatusTooManyRequests || !isOpenAIOAuthAccount(account) {
  2074. return false
  2075. }
  2076. if !isCodexSparkModel(requestedModel) || !account.ShouldHandleErrorCode(statusCode) {
  2077. return false
  2078. }
  2079. modelKey := normalizeCodexModel(modelRateLimitKeyForUpstreamModelNotFound(ctx, account, requestedModel))
  2080. if modelKey == "" {
  2081. return false
  2082. }
  2083. now := time.Now()
  2084. disposition, resetAt := classifyOpenAIOAuth429(headers, responseBody)
  2085. // Spark 只有明确耗尽 5h/7d 窗口时才能使用上游长 reset;普通瞬时 429
  2086. // 即使携带全局 reset 头,也只能使用短时回避,避免错误冷却数天。
  2087. if disposition != openAIOAuth429Quota5h && disposition != openAIOAuth429Quota7d {
  2088. resetAt = nil
  2089. }
  2090. if resetAt == nil || !resetAt.After(now) {
  2091. cooldown, ok := s.get429FallbackCooldown(ctx, account)
  2092. if !ok || cooldown <= 0 {
  2093. cooldown = time.Duration(defaultRateLimit429CooldownSeconds) * time.Second
  2094. }
  2095. reset := now.Add(cooldown)
  2096. resetAt = &reset
  2097. }
  2098. if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, modelKey, *resetAt, openAICodexSparkRateLimitReason); err != nil {
  2099. slog.Warn("openai_codex_spark_model_rate_limit_set_failed", "account_id", account.ID, "model", modelKey, "error", err)
  2100. }
  2101. slog.Info("openai_codex_spark_model_rate_limited", "account_id", account.ID, "model", modelKey, "reset_at", *resetAt)
  2102. return true
  2103. }
  2104. func (s *RateLimitService) HandleOpenAIImageCapabilityLoss(ctx context.Context, account *Account, statusCode int, responseBody []byte) bool {
  2105. if s == nil || account == nil || s.accountRepo == nil {
  2106. return false
  2107. }
  2108. if account.Platform != PlatformOpenAI {
  2109. return false
  2110. }
  2111. if !account.ShouldHandleErrorCode(statusCode) {
  2112. slog.Info("openai_image_capability_loss_skipped_by_error_code_policy", "account_id", account.ID, "status_code", statusCode)
  2113. return false
  2114. }
  2115. if !isOpenAIImageCapabilityLossError(statusCode, responseBody) {
  2116. return false
  2117. }
  2118. resetAt := time.Now().Add(openAIImageCapabilityLossCooldown)
  2119. if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, openAIImageGenerationRateLimitKey, resetAt, openAIImageCapabilityLossReason); err != nil {
  2120. slog.Warn("openai_image_capability_loss_set_model_rate_limit_failed", "account_id", account.ID, "scope", openAIImageGenerationRateLimitKey, "error", err)
  2121. return true
  2122. }
  2123. slog.Info("openai_image_capability_lost", "account_id", account.ID, "scope", openAIImageGenerationRateLimitKey, "reset_at", resetAt, "reset_in", time.Until(resetAt).Truncate(time.Second))
  2124. return true
  2125. }
  2126. // isOpenAIImageCapabilityLossError reports whether upstream rejected the
  2127. // image_generation tool choice that sub2api itself put into the request body.
  2128. // Only meaningful for self-built images requests, where tools always carries a
  2129. // matching image_generation entry — upstream saying otherwise means the account
  2130. // lost the capability.
  2131. func isOpenAIImageCapabilityLossError(statusCode int, body []byte) bool {
  2132. if statusCode != http.StatusBadRequest || len(body) == 0 {
  2133. return false
  2134. }
  2135. lower := strings.ToLower(string(body))
  2136. return strings.Contains(lower, "image_generation") &&
  2137. strings.Contains(lower, "not found in 'tools' parameter")
  2138. }
  2139. func isOpenAIImageRateLimitError(statusCode int, body []byte) bool {
  2140. if statusCode != http.StatusTooManyRequests || len(body) == 0 {
  2141. return false
  2142. }
  2143. lower := strings.ToLower(string(body))
  2144. for _, marker := range []string{
  2145. "for limit gpt-image",
  2146. "input-images per min",
  2147. "gpt-image-2-codex",
  2148. "gpt-image",
  2149. } {
  2150. if strings.Contains(lower, marker) {
  2151. return true
  2152. }
  2153. }
  2154. return false
  2155. }
  2156. func openAIImageRateLimitResetAt(headers http.Header, body []byte) time.Time {
  2157. now := time.Now()
  2158. if resetAt := parseRetryAfterResetTime(headers, now); resetAt != nil && resetAt.After(now) {
  2159. return *resetAt
  2160. }
  2161. if resetAt := calculateOpenAI429ResetTime(headers); resetAt != nil && resetAt.After(now) {
  2162. return *resetAt
  2163. }
  2164. if resetUnix := parseOpenAIRateLimitResetTime(body); resetUnix != nil {
  2165. if resetAt := time.Unix(*resetUnix, 0); resetAt.After(now) {
  2166. return resetAt
  2167. }
  2168. }
  2169. if cooldown := parseOpenAIImageTryAgainCooldown(body); cooldown > 0 {
  2170. return now.Add(cooldown)
  2171. }
  2172. return now.Add(openAIImageRateLimitDefaultCooldown)
  2173. }
  2174. func parseRetryAfterResetTime(headers http.Header, now time.Time) *time.Time {
  2175. if headers == nil {
  2176. return nil
  2177. }
  2178. raw := strings.TrimSpace(headers.Get("Retry-After"))
  2179. if raw == "" {
  2180. return nil
  2181. }
  2182. if seconds, err := strconv.ParseFloat(raw, 64); err == nil {
  2183. resetAt := now.Add(time.Duration(seconds * float64(time.Second)))
  2184. return &resetAt
  2185. }
  2186. if parsed, err := http.ParseTime(raw); err == nil {
  2187. return &parsed
  2188. }
  2189. return nil
  2190. }
  2191. func parseOpenAIImageTryAgainCooldown(body []byte) time.Duration {
  2192. if len(body) == 0 {
  2193. return 0
  2194. }
  2195. match := openAIImageTryAgainPattern.FindSubmatch(body)
  2196. if len(match) != 3 {
  2197. return 0
  2198. }
  2199. value, err := strconv.ParseFloat(string(match[1]), 64)
  2200. if err != nil || value <= 0 {
  2201. return 0
  2202. }
  2203. switch strings.ToLower(string(match[2])) {
  2204. case "ms":
  2205. return time.Duration(value * float64(time.Millisecond))
  2206. case "s", "sec", "secs", "second", "seconds":
  2207. return time.Duration(value * float64(time.Second))
  2208. case "m", "min", "mins", "minute", "minutes":
  2209. return time.Duration(value * float64(time.Minute))
  2210. default:
  2211. return 0
  2212. }
  2213. }
  2214. const upstreamModelNotFoundCooldown = 30 * time.Minute
  2215. const upstreamModelNotFoundReason = "upstream_404_model_not_found"
  2216. const upstreamCodexPlanGatedModelCooldown = 30 * time.Minute
  2217. const upstreamCodexPlanGatedModelReason = "upstream_400_codex_plan_gated_model"
  2218. const tempUnschedBodyMaxBytes = 64 << 10
  2219. const tempUnschedMessageMaxBytes = 2048
  2220. // HandleUpstreamModelNotFound marks the requested model as temporarily
  2221. // unavailable on the account when the upstream deterministically reports it
  2222. // cannot serve that model: a 404 model-not-found, or the Codex 400 rejecting a
  2223. // plan-gated model on a ChatGPT OAuth account. Returning true tells the caller
  2224. // to fail the current attempt over to another account; the scheduler skips the
  2225. // (account, model) pair via IsSchedulableForModelWithContext until the
  2226. // cooldown expires, instead of re-selecting an account that can never serve
  2227. // the model.
  2228. func (s *RateLimitService) HandleUpstreamModelNotFound(ctx context.Context, account *Account, requestedModel string, statusCode int, responseBody []byte) bool {
  2229. if s == nil || account == nil || s.accountRepo == nil {
  2230. return false
  2231. }
  2232. if !account.ShouldHandleErrorCode(statusCode) {
  2233. return false
  2234. }
  2235. var cooldown time.Duration
  2236. var reason string
  2237. switch {
  2238. case isUpstreamModelNotFoundError(statusCode, responseBody):
  2239. cooldown, reason = upstreamModelNotFoundCooldown, upstreamModelNotFoundReason
  2240. case isOpenAIOAuthAccount(account) && isOpenAICodexPlanGatedModelError(statusCode, responseBody):
  2241. cooldown, reason = upstreamCodexPlanGatedModelCooldown, upstreamCodexPlanGatedModelReason
  2242. default:
  2243. return false
  2244. }
  2245. modelKey := modelRateLimitKeyForUpstreamModelNotFound(ctx, account, requestedModel)
  2246. if modelKey == "" {
  2247. return false
  2248. }
  2249. if shouldSkipCodexPlanGatedImageModelCooldown(ctx, reason, requestedModel, modelKey) {
  2250. return true
  2251. }
  2252. resetAt := time.Now().Add(cooldown)
  2253. if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, modelKey, resetAt, reason); err != nil {
  2254. slog.Warn("upstream_model_not_found_set_model_rate_limit_failed", "account_id", account.ID, "model", modelKey, "reason", reason, "error", err)
  2255. return true
  2256. }
  2257. slog.Info("upstream_model_not_found_model_rate_limited", "account_id", account.ID, "model", modelKey, "reason", reason, "reset_at", resetAt)
  2258. return true
  2259. }
  2260. // shouldSkipCodexPlanGatedImageModelCooldown 判断这次 Codex plan-gated 400 是否
  2261. // 属于"图片模型被文本端点拒绝"。
  2262. //
  2263. // 这类错误是确定性的端点错配,不是账号能力缺失:同一账号通过 /v1/images/* 依然
  2264. // 能出图。在这里写 per-model 冷却,会让一次用错端点的请求把整个号池对正确的生图
  2265. // 端点下线(#4828)。
  2266. //
  2267. // 但请求本身就从 /v1/images/* 入站时不适用——那种情况下被拒说明账号确实不具备
  2268. // 该模型能力,冷却是必要的刹车:没有它,每个生图请求都会完整走一遍号池,对上游
  2269. // 形成无上界的 400 放大。
  2270. //
  2271. // 请求模型与最终冷却键都要判:冷却键走的是 account.GetMappedModel,账号可能把
  2272. // 文本别名映射到 gpt-image-*,只判请求模型会漏掉这种形态。
  2273. func shouldSkipCodexPlanGatedImageModelCooldown(ctx context.Context, reason, requestedModel, modelKey string) bool {
  2274. if reason != upstreamCodexPlanGatedModelReason {
  2275. return false
  2276. }
  2277. if OpenAIImagesEndpointFromContext(ctx) {
  2278. return false
  2279. }
  2280. return IsGPTImageGenerationModel(requestedModel) || IsGPTImageGenerationModel(modelKey)
  2281. }
  2282. func modelRateLimitKeyForUpstreamModelNotFound(ctx context.Context, account *Account, requestedModel string) string {
  2283. modelKey := strings.TrimSpace(requestedModel)
  2284. if account == nil || modelKey == "" {
  2285. return modelKey
  2286. }
  2287. if account.Platform == PlatformAntigravity {
  2288. if resolved := strings.TrimSpace(resolveFinalAntigravityModelKey(ctx, account, modelKey)); resolved != "" {
  2289. return resolved
  2290. }
  2291. return modelKey
  2292. }
  2293. if mapped := strings.TrimSpace(account.GetMappedModel(modelKey)); mapped != "" {
  2294. return mapped
  2295. }
  2296. return modelKey
  2297. }
  2298. func firstRequestedModel(requestedModel []string) string {
  2299. if len(requestedModel) == 0 {
  2300. return ""
  2301. }
  2302. return strings.TrimSpace(requestedModel[0])
  2303. }
  2304. type tempUnschedulableModelContextKey struct{}
  2305. func withTempUnschedulableModel(ctx context.Context, requestedModel []string) context.Context {
  2306. model := firstRequestedModel(requestedModel)
  2307. if model == "" {
  2308. return ctx
  2309. }
  2310. if ctx == nil {
  2311. ctx = context.Background()
  2312. }
  2313. return context.WithValue(ctx, tempUnschedulableModelContextKey{}, model)
  2314. }
  2315. func tempUnschedulableModel(ctx context.Context, requestedModel []string) string {
  2316. if model := firstRequestedModel(requestedModel); model != "" {
  2317. return model
  2318. }
  2319. if ctx == nil {
  2320. return ""
  2321. }
  2322. model, _ := ctx.Value(tempUnschedulableModelContextKey{}).(string)
  2323. return strings.TrimSpace(model)
  2324. }
  2325. type tempUnschedulableRuleMatch struct {
  2326. rule TempUnschedulableRule
  2327. ruleIndex int
  2328. matchedKeyword string
  2329. }
  2330. func matchTempUnschedulableRules(account *Account, statusCode int, responseBody []byte) []tempUnschedulableRuleMatch {
  2331. if account == nil || !account.IsTempUnschedulableEnabled() || statusCode <= 0 || len(responseBody) == 0 {
  2332. return nil
  2333. }
  2334. rules := account.GetTempUnschedulableRules()
  2335. if len(rules) == 0 {
  2336. return nil
  2337. }
  2338. body := responseBody
  2339. if len(body) > tempUnschedBodyMaxBytes {
  2340. body = body[:tempUnschedBodyMaxBytes]
  2341. }
  2342. bodyLower := strings.ToLower(string(body))
  2343. matches := make([]tempUnschedulableRuleMatch, 0, 1)
  2344. for idx, rule := range rules {
  2345. if rule.ErrorCode != statusCode || len(rule.Keywords) == 0 {
  2346. continue
  2347. }
  2348. matchedKeyword := matchTempUnschedKeyword(bodyLower, rule.Keywords)
  2349. if matchedKeyword == "" {
  2350. continue
  2351. }
  2352. matches = append(matches, tempUnschedulableRuleMatch{rule: rule, ruleIndex: idx, matchedKeyword: matchedKeyword})
  2353. }
  2354. return matches
  2355. }
  2356. func (s *RateLimitService) tryTempUnschedulable(ctx context.Context, account *Account, statusCode int, responseBody []byte, requestedModel ...string) bool {
  2357. if account == nil {
  2358. return false
  2359. }
  2360. if !account.IsTempUnschedulableEnabled() {
  2361. return false
  2362. }
  2363. // 401 首次命中可临时不可调度(给 token 刷新窗口);
  2364. // 若历史上已因 401 进入过临时不可调度,则本次应升级为 error(返回 false 交由默认错误逻辑处理)。
  2365. // Antigravity 跳过:其 401 由 applyErrorPolicy 的 temp_unschedulable_rules 自行控制,无需升级逻辑。
  2366. if statusCode == http.StatusUnauthorized && account.Platform != PlatformAntigravity {
  2367. reason := account.TempUnschedulableReason
  2368. // 缓存可能没有 reason,从 DB 回退读取
  2369. if reason == "" {
  2370. if dbAcc, err := s.accountRepo.GetByID(ctx, account.ID); err == nil && dbAcc != nil {
  2371. reason = dbAcc.TempUnschedulableReason
  2372. }
  2373. }
  2374. if wasTempUnschedByStatusCode(reason, statusCode) {
  2375. slog.Info("401_escalated_to_error", "account_id", account.ID,
  2376. "reason", "previous temp-unschedulable was also 401")
  2377. return false
  2378. }
  2379. }
  2380. for _, match := range matchTempUnschedulableRules(account, statusCode, responseBody) {
  2381. if s.triggerTempUnschedulable(ctx, account, match.rule, match.ruleIndex, statusCode, match.matchedKeyword, responseBody, tempUnschedulableModel(ctx, requestedModel)) {
  2382. return true
  2383. }
  2384. }
  2385. return false
  2386. }
  2387. func wasTempUnschedByStatusCode(reason string, statusCode int) bool {
  2388. if statusCode <= 0 {
  2389. return false
  2390. }
  2391. reason = strings.TrimSpace(reason)
  2392. if reason == "" {
  2393. return false
  2394. }
  2395. var state TempUnschedState
  2396. if err := json.Unmarshal([]byte(reason), &state); err != nil {
  2397. return false
  2398. }
  2399. return state.StatusCode == statusCode
  2400. }
  2401. func matchTempUnschedKeyword(bodyLower string, keywords []string) string {
  2402. if bodyLower == "" {
  2403. return ""
  2404. }
  2405. for _, keyword := range keywords {
  2406. k := strings.TrimSpace(keyword)
  2407. if k == "" {
  2408. continue
  2409. }
  2410. if strings.Contains(bodyLower, strings.ToLower(k)) {
  2411. return k
  2412. }
  2413. }
  2414. return ""
  2415. }
  2416. func (s *RateLimitService) triggerTempUnschedulable(ctx context.Context, account *Account, rule TempUnschedulableRule, ruleIndex int, statusCode int, matchedKeyword string, responseBody []byte, requestedModel ...string) bool {
  2417. if account == nil {
  2418. return false
  2419. }
  2420. if rule.DurationMinutes <= 0 {
  2421. return false
  2422. }
  2423. now := time.Now()
  2424. until := now.Add(time.Duration(rule.DurationMinutes) * time.Minute)
  2425. state := &TempUnschedState{
  2426. UntilUnix: until.Unix(),
  2427. TriggeredAtUnix: now.Unix(),
  2428. StatusCode: statusCode,
  2429. MatchedKeyword: matchedKeyword,
  2430. RuleIndex: ruleIndex,
  2431. ErrorMessage: truncateTempUnschedMessage(responseBody, tempUnschedMessageMaxBytes),
  2432. }
  2433. reason := ""
  2434. if raw, err := json.Marshal(state); err == nil {
  2435. reason = string(raw)
  2436. }
  2437. if reason == "" {
  2438. reason = strings.TrimSpace(state.ErrorMessage)
  2439. }
  2440. // Persist known-model failures under the model key so the scheduler excludes
  2441. // only this (account, model) pair. Authentication and model-unknown failures
  2442. // retain the legacy account-wide temporary-unschedulable behavior below.
  2443. modelKey := firstRequestedModel(requestedModel)
  2444. if modelKey != "" && statusCode != http.StatusUnauthorized {
  2445. if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, modelKey, until, reason); err != nil {
  2446. slog.Warn("temp_unsched_model_rate_limit_set_failed", "account_id", account.ID, "model", modelKey, "error", err)
  2447. // The rule matched, so fail over the current request even if persistence
  2448. // failed; never widen a model-scoped failure into an account-wide block.
  2449. return true
  2450. }
  2451. slog.Info("account_model_temp_unschedulable", "account_id", account.ID, "model", modelKey, "until", until, "rule_index", ruleIndex, "status_code", statusCode)
  2452. return true
  2453. }
  2454. s.notifyAccountSchedulingBlocked(account, until, "temp_unschedulable")
  2455. if err := s.accountRepo.SetTempUnschedulable(ctx, account.ID, until, reason); err != nil {
  2456. slog.Warn("temp_unsched_set_failed", "account_id", account.ID, "error", err)
  2457. return false
  2458. }
  2459. if s.tempUnschedCache != nil {
  2460. if err := s.tempUnschedCache.SetTempUnsched(ctx, account.ID, state); err != nil {
  2461. slog.Warn("temp_unsched_cache_set_failed", "account_id", account.ID, "error", err)
  2462. }
  2463. }
  2464. slog.Info("account_temp_unschedulable", "account_id", account.ID, "until", until, "rule_index", ruleIndex, "status_code", statusCode)
  2465. return true
  2466. }
  2467. func truncateTempUnschedMessage(body []byte, maxBytes int) string {
  2468. if maxBytes <= 0 || len(body) == 0 {
  2469. return ""
  2470. }
  2471. if len(body) > maxBytes {
  2472. body = body[:maxBytes]
  2473. }
  2474. return strings.TrimSpace(string(body))
  2475. }
  2476. // HandleStreamTimeout 处理流数据超时
  2477. // 根据系统设置决定是否标记账户为临时不可调度或错误状态
  2478. // 返回是否应该停止该账号的调度
  2479. func (s *RateLimitService) HandleStreamTimeout(ctx context.Context, account *Account, model string) bool {
  2480. if account == nil {
  2481. return false
  2482. }
  2483. // 获取系统设置
  2484. if s.settingService == nil {
  2485. slog.Warn("stream_timeout_setting_service_missing", "account_id", account.ID)
  2486. return false
  2487. }
  2488. settings, err := s.settingService.GetStreamTimeoutSettings(ctx)
  2489. if err != nil {
  2490. slog.Warn("stream_timeout_get_settings_failed", "account_id", account.ID, "error", err)
  2491. return false
  2492. }
  2493. if !settings.Enabled {
  2494. return false
  2495. }
  2496. if settings.Action == StreamTimeoutActionNone {
  2497. return false
  2498. }
  2499. // 增加超时计数
  2500. var count int64 = 1
  2501. if s.timeoutCounterCache != nil {
  2502. count, err = s.timeoutCounterCache.IncrementTimeoutCount(ctx, account.ID, settings.ThresholdWindowMinutes)
  2503. if err != nil {
  2504. slog.Warn("stream_timeout_increment_count_failed", "account_id", account.ID, "error", err)
  2505. // 继续处理,使用 count=1
  2506. count = 1
  2507. }
  2508. }
  2509. slog.Info("stream_timeout_count", "account_id", account.ID, "count", count, "threshold", settings.ThresholdCount, "window_minutes", settings.ThresholdWindowMinutes, "model", model)
  2510. // 检查是否达到阈值
  2511. if count < int64(settings.ThresholdCount) {
  2512. return false
  2513. }
  2514. // 达到阈值,执行相应操作
  2515. switch settings.Action {
  2516. case StreamTimeoutActionTempUnsched:
  2517. return s.triggerStreamTimeoutTempUnsched(ctx, account, settings, model)
  2518. case StreamTimeoutActionError:
  2519. return s.triggerStreamTimeoutError(ctx, account, model)
  2520. default:
  2521. return false
  2522. }
  2523. }
  2524. // triggerStreamTimeoutTempUnsched 触发流超时临时不可调度
  2525. func (s *RateLimitService) triggerStreamTimeoutTempUnsched(ctx context.Context, account *Account, settings *StreamTimeoutSettings, model string) bool {
  2526. now := time.Now()
  2527. until := now.Add(time.Duration(settings.TempUnschedMinutes) * time.Minute)
  2528. state := &TempUnschedState{
  2529. UntilUnix: until.Unix(),
  2530. TriggeredAtUnix: now.Unix(),
  2531. StatusCode: 0, // 超时没有状态码
  2532. MatchedKeyword: "stream_timeout",
  2533. RuleIndex: -1, // 表示系统级规则
  2534. ErrorMessage: "Stream data interval timeout for model: " + model,
  2535. }
  2536. reason := ""
  2537. if raw, err := json.Marshal(state); err == nil {
  2538. reason = string(raw)
  2539. }
  2540. if reason == "" {
  2541. reason = state.ErrorMessage
  2542. }
  2543. s.notifyAccountSchedulingBlocked(account, until, "stream_timeout_temp_unschedulable")
  2544. if err := s.accountRepo.SetTempUnschedulable(ctx, account.ID, until, reason); err != nil {
  2545. slog.Warn("stream_timeout_set_temp_unsched_failed", "account_id", account.ID, "error", err)
  2546. return false
  2547. }
  2548. if s.tempUnschedCache != nil {
  2549. if err := s.tempUnschedCache.SetTempUnsched(ctx, account.ID, state); err != nil {
  2550. slog.Warn("stream_timeout_set_temp_unsched_cache_failed", "account_id", account.ID, "error", err)
  2551. }
  2552. }
  2553. // 重置超时计数
  2554. if s.timeoutCounterCache != nil {
  2555. if err := s.timeoutCounterCache.ResetTimeoutCount(ctx, account.ID); err != nil {
  2556. slog.Warn("stream_timeout_reset_count_failed", "account_id", account.ID, "error", err)
  2557. }
  2558. }
  2559. slog.Info("stream_timeout_temp_unschedulable", "account_id", account.ID, "until", until, "model", model)
  2560. return true
  2561. }
  2562. // triggerStreamTimeoutError 触发流超时错误状态
  2563. func (s *RateLimitService) triggerStreamTimeoutError(ctx context.Context, account *Account, model string) bool {
  2564. errorMsg := "Stream data interval timeout (repeated failures) for model: " + model
  2565. s.notifyAccountSchedulingBlocked(account, time.Time{}, "stream_timeout_error")
  2566. if err := s.accountRepo.SetError(ctx, account.ID, errorMsg); err != nil {
  2567. slog.Warn("stream_timeout_set_error_failed", "account_id", account.ID, "error", err)
  2568. return false
  2569. }
  2570. // 重置超时计数
  2571. if s.timeoutCounterCache != nil {
  2572. if err := s.timeoutCounterCache.ResetTimeoutCount(ctx, account.ID); err != nil {
  2573. slog.Warn("stream_timeout_reset_count_failed", "account_id", account.ID, "error", err)
  2574. }
  2575. }
  2576. slog.Warn("stream_timeout_account_error", "account_id", account.ID, "model", model)
  2577. return true
  2578. }