openai_account_runtime_block_fastpath.go 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611
  1. package service
  2. import (
  3. "context"
  4. "log/slog"
  5. "net/http"
  6. "strings"
  7. "sync"
  8. "time"
  9. )
  10. const (
  11. openAIAccountStateUpdateTimeout = 5 * time.Second
  12. openAIOAuth429FallbackCooldown = 5 * time.Second
  13. openAIOAuth429RetryWindow = 2 * time.Minute
  14. openAIOAuth429RetryDelay = time.Second
  15. openAIOAuth429MaxSameAccountRetries = 3
  16. openAIOAuth429MaxAccountAttempts = 3
  17. openAIStopSchedulingBridgeCooldown = 2 * time.Minute
  18. openAIOAuth429StormWindow = 10 * time.Second
  19. openAIOAuth429StormMaxAccountSwitches = 1
  20. )
  21. // OpenAIOAuth429FailoverState tracks the request-local follow-up budget after
  22. // the first Grok OAuth 429. Once that 429 occurs, exactly one different account
  23. // may be attempted; any failure from that follow-up account ends failover.
  24. type OpenAIOAuth429FailoverState struct {
  25. grokOAuth429FollowupPending bool
  26. }
  27. type openAIOAuth429Disposition uint8
  28. const (
  29. openAIOAuth429Transient openAIOAuth429Disposition = iota
  30. openAIOAuth429Quota5h
  31. openAIOAuth429Quota7d
  32. openAIOAuth429QuotaReset
  33. openAIOAuth429RetryAfter
  34. )
  35. // classifyOpenAIOAuth429 区分账号配额耗尽信号与普通瞬时 429。只有窗口达到
  36. // 100% 或响应体明确给出 reset 时间时,才视为配额限流信号。
  37. func classifyOpenAIOAuth429(headers http.Header, responseBody []byte) (openAIOAuth429Disposition, *time.Time) {
  38. if snapshot := ParseCodexRateLimitHeaders(headers); snapshot != nil {
  39. if normalized := snapshot.Normalize(); normalized != nil {
  40. if normalized.Used7dPercent != nil && *normalized.Used7dPercent >= 100 {
  41. if normalized.Reset7dSeconds != nil {
  42. now := time.Now()
  43. resetAt := now.Add(time.Duration(*normalized.Reset7dSeconds) * time.Second)
  44. return openAIOAuth429Quota7d, &resetAt
  45. }
  46. return openAIOAuth429Quota7d, nil
  47. }
  48. if normalized.Used5hPercent != nil && *normalized.Used5hPercent >= 100 {
  49. if normalized.Reset5hSeconds != nil {
  50. now := time.Now()
  51. resetAt := now.Add(time.Duration(*normalized.Reset5hSeconds) * time.Second)
  52. return openAIOAuth429Quota5h, &resetAt
  53. }
  54. return openAIOAuth429Quota5h, nil
  55. }
  56. }
  57. }
  58. if resetAt := calculateOpenAI429ResetTime(headers); resetAt != nil {
  59. return openAIOAuth429QuotaReset, resetAt
  60. }
  61. if resetUnix := parseOpenAIRateLimitResetTime(responseBody); resetUnix != nil {
  62. resetAt := time.Unix(*resetUnix, 0)
  63. return openAIOAuth429QuotaReset, &resetAt
  64. }
  65. if resetAt := parseRetryAfterResetTime(headers, time.Now()); resetAt != nil && resetAt.After(time.Now()) {
  66. return openAIOAuth429RetryAfter, resetAt
  67. }
  68. return openAIOAuth429Transient, nil
  69. }
  70. func openAIAccountStateContext(ctx context.Context) (context.Context, context.CancelFunc) {
  71. base := context.Background()
  72. if ctx != nil {
  73. base = context.WithoutCancel(ctx)
  74. }
  75. return context.WithTimeout(base, openAIAccountStateUpdateTimeout)
  76. }
  77. func isOpenAIOAuthAccount(account *Account) bool {
  78. return account != nil && account.IsOpenAIOAuthLike()
  79. }
  80. func isGrokOAuthAccount(account *Account) bool {
  81. return account != nil && account.Platform == PlatformGrok && account.Type == AccountTypeOAuth
  82. }
  83. func isOpenAIAccount(account *Account) bool {
  84. return account != nil && (account.Platform == PlatformOpenAI || account.Platform == PlatformGrok)
  85. }
  86. // handleOpenAIAccountUpstreamError expects canonicalModel to be the model used
  87. // for scheduling after applying account mapping exactly once.
  88. func (s *OpenAIGatewayService) handleOpenAIAccountUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte, canonicalModel ...string) bool {
  89. if account != nil && account.Platform == PlatformGrok && isGrokContentPolicyRejection(statusCode, responseBody) {
  90. return false
  91. }
  92. // Any non-2xx upstream HTTP response means the model request was actually sent.
  93. if s != nil {
  94. scheduleOllamaCloudUsageActivity(s.deferredService, account)
  95. }
  96. // Capacity shedding describes this request, not account health. Keep the
  97. // account schedulable while the request-local retry budget handles recovery.
  98. if account != nil && account.Platform == PlatformOpenAI && isOpenAIRequestScopedCapacityShed("", responseBody) {
  99. return false
  100. }
  101. stateCtx, cancel := openAIAccountStateContext(ctx)
  102. defer cancel()
  103. if account != nil && account.Platform == PlatformOpenAI && isOpenAIHTTPUpstreamAccessStateError(statusCode, "", responseBody) {
  104. message := "OpenAI upstream account or workspace is unavailable"
  105. if upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(responseBody)); upstreamMsg != "" {
  106. message = upstreamMsg
  107. }
  108. if s != nil && s.rateLimitService != nil {
  109. s.rateLimitService.handleAuthError(stateCtx, account, message)
  110. }
  111. if s != nil {
  112. s.BlockAccountScheduling(account, time.Time{}, "openai_access_state")
  113. }
  114. return true
  115. }
  116. if account != nil && account.Platform == PlatformOpenAI && isOpenAIContextWindowError("", responseBody) {
  117. return false
  118. }
  119. if isOpenAIImageRateLimitError(statusCode, responseBody) {
  120. if s != nil && s.rateLimitService != nil {
  121. _ = s.rateLimitService.HandleOpenAIImageRateLimit(stateCtx, account, statusCode, headers, responseBody)
  122. }
  123. return false
  124. }
  125. // Self-built images requests always carry a matching image_generation tool, so a
  126. // "tool choice not found in 'tools'" 400 means upstream revoked this account's
  127. // image capability. Gated on the self-built marker: passthrough clients control
  128. // their own tools/tool_choice and could otherwise poison a healthy account.
  129. if isOpenAIImagesSelfBuiltRequest(ctx) && isOpenAIImageCapabilityLossError(statusCode, responseBody) {
  130. if s != nil && s.rateLimitService != nil {
  131. _ = s.rateLimitService.HandleOpenAIImageCapabilityLoss(stateCtx, account, statusCode, responseBody)
  132. }
  133. return false
  134. }
  135. if s == nil || account == nil {
  136. return false
  137. }
  138. // Team 联动熔断必须先于 model-not-found 与账户级临时不可调度规则的早退。
  139. if s.rateLimitService != nil {
  140. s.rateLimitService.maybeHandleOpenAITeamLinkedError(stateCtx, account, statusCode, responseBody)
  141. }
  142. stateCtx = withTempUnschedulableModel(stateCtx, canonicalModel)
  143. if s.rateLimitService != nil && len(canonicalModel) > 0 && s.rateLimitService.HandleUpstreamModelNotFound(stateCtx, account, canonicalModel[0], statusCode, responseBody) {
  144. return true
  145. }
  146. // Isolate a custom temporary-unschedulable match to the known upstream
  147. // model before entering the generic account error path. This keeps the
  148. // account available to other models and avoids the account runtime blocker.
  149. if s.rateLimitService != nil && statusCode != http.StatusUnauthorized && len(canonicalModel) > 0 && strings.TrimSpace(canonicalModel[0]) != "" &&
  150. s.rateLimitService.HandleTempUnschedulable(stateCtx, account, statusCode, responseBody, canonicalModel[0]) {
  151. return true
  152. }
  153. if statusCode == http.StatusTooManyRequests && s.rateLimitService != nil && len(canonicalModel) > 0 &&
  154. s.rateLimitService.HandleOpenAICodexSparkRateLimit(stateCtx, account, canonicalModel[0], statusCode, headers, responseBody) {
  155. return false
  156. }
  157. if statusCode == http.StatusTooManyRequests {
  158. s.markOpenAIOAuth429RateLimited(stateCtx, account, headers, responseBody)
  159. }
  160. if s.rateLimitService == nil {
  161. return false
  162. }
  163. shouldDisable := s.rateLimitService.HandleUpstreamError(stateCtx, account, statusCode, headers, responseBody)
  164. modelTempMatched := statusCode != http.StatusUnauthorized && tempUnschedulableModel(stateCtx, nil) != "" &&
  165. len(matchTempUnschedulableRules(account, statusCode, responseBody)) > 0
  166. if shouldDisable && !modelTempMatched {
  167. s.BlockAccountScheduling(account, time.Time{}, "upstream_disable")
  168. }
  169. // Pool-mode retryable upstream errors are already bounded by the request-local
  170. // same-account retry budget. Recording the generic account+model transient
  171. // cooldown here would block the next approved retry before that budget is used.
  172. poolModeRetryable := account.IsPoolMode() && account.IsPoolModeRetryableStatus(statusCode)
  173. if !shouldDisable && account.Platform == PlatformOpenAI && account.Type == AccountTypeAPIKey &&
  174. shouldCooldownOpenAITransientUpstreamError(statusCode, responseBody) && !poolModeRetryable {
  175. model := ""
  176. if len(canonicalModel) > 0 {
  177. model = canonicalModel[0]
  178. }
  179. decision := s.recordOpenAIAccountModelTransientFailure(account, model, time.Now())
  180. if decision.FailureStreak > 0 {
  181. slog.Warn("openai_model_transient_state",
  182. "account_id", account.ID,
  183. "model", openAIAccountModelTransientModel(model),
  184. "failure_streak", decision.FailureStreak,
  185. "cooldown_ms", decision.Cooldown.Milliseconds(),
  186. "block_scope", "account_model",
  187. )
  188. }
  189. }
  190. return shouldDisable
  191. }
  192. func shouldCooldownOpenAITransientUpstreamError(statusCode int, responseBody []byte) bool {
  193. switch statusCode {
  194. case http.StatusInternalServerError, http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout, 520, 521, 522, 523, 524:
  195. return true
  196. case http.StatusBadRequest:
  197. return isOpenAITransientProcessingError(statusCode, "", responseBody)
  198. default:
  199. return false
  200. }
  201. }
  202. func (s *OpenAIGatewayService) markOpenAIOAuth429RateLimited(ctx context.Context, account *Account, headers http.Header, responseBody []byte) {
  203. if s == nil || !isOpenAIOAuthAccount(account) {
  204. return
  205. }
  206. // Spark 影子:不按 /responses 429 的 global x-codex-* 信号做内存运行时熔断(同 handle429,外审第8轮 P1)。
  207. // 同时避免把 spark 的 429 计入全局 429 storm 计数(recordOpenAIOAuth429),否则会误伤母账号 failover 决策。
  208. if account.IsShadow() {
  209. return
  210. }
  211. s.recordOpenAIOAuth429()
  212. disposition, resetAt := classifyOpenAIOAuth429(headers, responseBody)
  213. if disposition == openAIOAuth429Transient && s.openAIOAuth429RetryWindowActive(account) {
  214. return
  215. }
  216. now := time.Now()
  217. cooldownUntil := now.Add(openAIOAuth429FallbackCooldown)
  218. if resetAt != nil && resetAt.After(now) {
  219. cooldownUntil = *resetAt
  220. } else if s.rateLimitService != nil {
  221. cooldown, ok := s.rateLimitService.get429FallbackCooldown(ctx, account)
  222. if !ok || cooldown <= 0 {
  223. s.openaiOAuth429RetryStartedAt.Delete(account.ID)
  224. return
  225. }
  226. cooldownUntil = now.Add(cooldown)
  227. }
  228. s.BlockAccountScheduling(account, cooldownUntil, "429")
  229. s.openaiOAuth429RetryStartedAt.Delete(account.ID)
  230. }
  231. func (s *OpenAIGatewayService) shouldRetryOpenAIOAuth429OnSameAccount(account *Account, statusCode int, shouldDisable bool) bool {
  232. return s.shouldRetryOpenAIOAuth429OnSameAccountWithResponse(account, statusCode, shouldDisable, nil, nil)
  233. }
  234. func (s *OpenAIGatewayService) shouldRetryOpenAIOAuth429OnSameAccountWithResponse(account *Account, statusCode int, shouldDisable bool, headers http.Header, responseBody []byte) bool {
  235. if shouldDisable || statusCode != http.StatusTooManyRequests || !isOpenAIOAuthAccount(account) || account.IsShadow() {
  236. return false
  237. }
  238. disposition, _ := classifyOpenAIOAuth429(headers, responseBody)
  239. if disposition != openAIOAuth429Transient {
  240. return false
  241. }
  242. // markOpenAIOAuth429RateLimited parks the account once the window expires.
  243. // Do not accidentally create a fresh window after that transition.
  244. if s.isOpenAIAccountRuntimeBlocked(account) {
  245. return false
  246. }
  247. return s.openAIOAuth429RetryWindowActive(account)
  248. }
  249. // ShouldRetryOpenAIOAuth429 lets RateLimitService defer persistent account
  250. // cooldown until the gateway's same-account retry window is exhausted.
  251. func (s *OpenAIGatewayService) ShouldRetryOpenAIOAuth429(account *Account, headers http.Header, responseBody []byte) bool {
  252. if s == nil || !isOpenAIOAuthAccount(account) || account.IsShadow() || s.isOpenAIAccountRuntimeBlocked(account) {
  253. return false
  254. }
  255. disposition, _ := classifyOpenAIOAuth429(headers, responseBody)
  256. if disposition != openAIOAuth429Transient {
  257. return false
  258. }
  259. return s.openAIOAuth429RetryWindowActive(account)
  260. }
  261. func (s *OpenAIGatewayService) openAIOAuth429RetryWindowActive(account *Account) bool {
  262. if s == nil || !isOpenAIOAuthAccount(account) || account.IsShadow() {
  263. return false
  264. }
  265. now := time.Now()
  266. value, _ := s.openaiOAuth429RetryStartedAt.LoadOrStore(account.ID, now)
  267. startedAt, ok := value.(time.Time)
  268. if !ok {
  269. s.openaiOAuth429RetryStartedAt.Store(account.ID, now)
  270. startedAt = now
  271. }
  272. return now.Before(startedAt.Add(openAIOAuth429RetryWindow))
  273. }
  274. func (s *OpenAIGatewayService) openAIOAuth429RetryDeadline(account *Account) time.Time {
  275. if s == nil || !isOpenAIOAuthAccount(account) || account.IsShadow() {
  276. return time.Time{}
  277. }
  278. value, ok := s.openaiOAuth429RetryStartedAt.Load(account.ID)
  279. if !ok {
  280. return time.Time{}
  281. }
  282. startedAt, ok := value.(time.Time)
  283. if !ok {
  284. return time.Time{}
  285. }
  286. return startedAt.Add(openAIOAuth429RetryWindow)
  287. }
  288. func openAIOAuth429SameAccountRetryDelay(headers http.Header, deadline time.Time) time.Duration {
  289. delay := openAIOAuth429RetryDelay
  290. now := time.Now()
  291. if resetAt := parseRetryAfterResetTime(headers, now); resetAt != nil && resetAt.After(now) {
  292. if wait := resetAt.Sub(now); wait > delay {
  293. delay = wait
  294. }
  295. }
  296. // A deadline limits whether a retry is allowed, never the upstream's wait.
  297. // The handler rejects attempts that cannot fit inside the remaining budget.
  298. return delay
  299. }
  300. func (s *OpenAIGatewayService) BlockAccountScheduling(account *Account, until time.Time, reason string) {
  301. if s == nil || !isOpenAIAccount(account) {
  302. return
  303. }
  304. mu := s.openAIAccountRuntimeBlockLock(account.ID)
  305. mu.Lock()
  306. defer mu.Unlock()
  307. _, _ = s.blockAccountSchedulingLocked(account, until, reason)
  308. }
  309. func (s *OpenAIGatewayService) openAIAccountRuntimeBlockLock(accountID int64) *sync.Mutex {
  310. actual, _ := s.openaiAccountRuntimeBlockLocks.LoadOrStore(accountID, &sync.Mutex{})
  311. mu, ok := actual.(*sync.Mutex)
  312. if !ok {
  313. mu = &sync.Mutex{}
  314. s.openaiAccountRuntimeBlockLocks.Store(accountID, mu)
  315. }
  316. return mu
  317. }
  318. func (s *OpenAIGatewayService) blockAccountSchedulingLocked(account *Account, until time.Time, _ string) (uint64, bool) {
  319. generation := s.openaiAccountRuntimeBlockSequence.Add(1)
  320. s.openaiAccountRuntimeBlockGeneration.Store(account.ID, generation)
  321. now := time.Now()
  322. blockUntil := until
  323. if blockUntil.IsZero() || !blockUntil.After(now) {
  324. blockUntil = now.Add(openAIStopSchedulingBridgeCooldown)
  325. }
  326. for {
  327. current, loaded := s.openaiAccountRuntimeBlockUntil.Load(account.ID)
  328. if !loaded {
  329. actual, stored := s.openaiAccountRuntimeBlockUntil.LoadOrStore(account.ID, blockUntil)
  330. if !stored {
  331. return generation, true
  332. }
  333. current = actual
  334. }
  335. currentUntil, ok := current.(time.Time)
  336. if !ok || currentUntil.IsZero() {
  337. if s.openaiAccountRuntimeBlockUntil.CompareAndSwap(account.ID, current, blockUntil) {
  338. return generation, true
  339. }
  340. continue
  341. }
  342. if !blockUntil.After(currentUntil) {
  343. return generation, false
  344. }
  345. if s.openaiAccountRuntimeBlockUntil.CompareAndSwap(account.ID, current, blockUntil) {
  346. return generation, true
  347. }
  348. }
  349. }
  350. func (s *OpenAIGatewayService) ClearAccountSchedulingBlock(accountID int64) {
  351. if s == nil || accountID <= 0 {
  352. return
  353. }
  354. mu := s.openAIAccountRuntimeBlockLock(accountID)
  355. mu.Lock()
  356. defer mu.Unlock()
  357. s.openaiAccountRuntimeBlockUntil.Delete(accountID)
  358. s.openaiOAuth429RetryStartedAt.Delete(accountID)
  359. s.openaiAccountRuntimeBlockGeneration.Store(accountID, s.openaiAccountRuntimeBlockSequence.Add(1))
  360. }
  361. func (s *OpenAIGatewayService) isOpenAIAccountRuntimeBlocked(account *Account) bool {
  362. if s == nil || !isOpenAIAccount(account) {
  363. return false
  364. }
  365. mu := s.openAIAccountRuntimeBlockLock(account.ID)
  366. mu.Lock()
  367. defer mu.Unlock()
  368. value, ok := s.openaiAccountRuntimeBlockUntil.Load(account.ID)
  369. if !ok {
  370. return false
  371. }
  372. cooldownUntil, ok := value.(time.Time)
  373. if !ok || cooldownUntil.IsZero() {
  374. s.openaiAccountRuntimeBlockUntil.Delete(account.ID)
  375. s.openaiAccountRuntimeBlockGeneration.Store(account.ID, s.openaiAccountRuntimeBlockSequence.Add(1))
  376. return false
  377. }
  378. if time.Now().Before(cooldownUntil) {
  379. return true
  380. }
  381. s.openaiAccountRuntimeBlockUntil.Delete(account.ID)
  382. s.openaiAccountRuntimeBlockGeneration.Store(account.ID, s.openaiAccountRuntimeBlockSequence.Add(1))
  383. return false
  384. }
  385. func (s *OpenAIGatewayService) getOpenAIAccountModelTransientState() *openAIAccountModelTransientState {
  386. if s == nil {
  387. return nil
  388. }
  389. s.openaiModelTransientOnce.Do(func() {
  390. if s.openaiModelTransient == nil {
  391. s.openaiModelTransient = newOpenAIAccountModelTransientState(openAIModelTransientDefaultMax)
  392. }
  393. })
  394. return s.openaiModelTransient
  395. }
  396. func canonicalOpenAIAccountSchedulingModel(account *Account, requestedModel string) string {
  397. model := strings.TrimSpace(requestedModel)
  398. if account == nil || model == "" {
  399. return model
  400. }
  401. if account.IsOpenAI() {
  402. return resolveOpenAIAccountUpstreamModelForRequest(account, model, false)
  403. }
  404. if mapped := strings.TrimSpace(account.GetMappedModel(model)); mapped != "" {
  405. return mapped
  406. }
  407. return model
  408. }
  409. func openAIAccountModelTransientModel(canonicalModel string) string {
  410. return normalizeOpenAIAccountModelTransientModel(canonicalModel)
  411. }
  412. func (s *OpenAIGatewayService) recordOpenAIAccountModelTransientFailure(account *Account, canonicalModel string, now time.Time) openAIAccountModelTransientDecision {
  413. if s == nil || account == nil {
  414. return openAIAccountModelTransientDecision{}
  415. }
  416. state := s.getOpenAIAccountModelTransientState()
  417. if state == nil {
  418. return openAIAccountModelTransientDecision{}
  419. }
  420. return state.recordFailure(account.ID, openAIAccountModelTransientModel(canonicalModel), now)
  421. }
  422. func (s *OpenAIGatewayService) clearOpenAIAccountModelTransientState(accountID int64, model string) {
  423. state := s.getOpenAIAccountModelTransientState()
  424. if state == nil {
  425. return
  426. }
  427. state.recordSuccess(accountID, model)
  428. }
  429. func (s *OpenAIGatewayService) isOpenAIAccountModelRuntimeBlocked(account *Account, requestedModel string) bool {
  430. if s == nil || account == nil {
  431. return false
  432. }
  433. state := s.getOpenAIAccountModelTransientState()
  434. if state == nil {
  435. return false
  436. }
  437. canonicalModel := canonicalOpenAIAccountSchedulingModel(account, requestedModel)
  438. return state.isBlocked(account.ID, openAIAccountModelTransientModel(canonicalModel), time.Now())
  439. }
  440. func accountPersistedSchedulingCooldownActive(account *Account) bool {
  441. if account == nil {
  442. return false
  443. }
  444. now := time.Now()
  445. if account.TempUnschedulableUntil != nil && now.Before(*account.TempUnschedulableUntil) {
  446. return true
  447. }
  448. if account.RateLimitResetAt != nil && now.Before(*account.RateLimitResetAt) {
  449. return true
  450. }
  451. if account.OverloadUntil != nil && now.Before(*account.OverloadUntil) {
  452. return true
  453. }
  454. return false
  455. }
  456. type openAIAccountRuntimeBlockSnapshot struct {
  457. until time.Time
  458. generation uint64
  459. blocked bool
  460. }
  461. func (s *OpenAIGatewayService) peekOpenAIAccountRuntimeBlock(account *Account) openAIAccountRuntimeBlockSnapshot {
  462. if s == nil || !isOpenAIAccount(account) {
  463. return openAIAccountRuntimeBlockSnapshot{}
  464. }
  465. mu := s.openAIAccountRuntimeBlockLock(account.ID)
  466. mu.Lock()
  467. defer mu.Unlock()
  468. value, ok := s.openaiAccountRuntimeBlockUntil.Load(account.ID)
  469. if !ok {
  470. return openAIAccountRuntimeBlockSnapshot{}
  471. }
  472. until, isTime := value.(time.Time)
  473. if !isTime || until.IsZero() || !time.Now().Before(until) {
  474. s.openaiAccountRuntimeBlockUntil.Delete(account.ID)
  475. s.openaiOAuth429RetryStartedAt.Delete(account.ID)
  476. s.openaiAccountRuntimeBlockGeneration.Store(account.ID, s.openaiAccountRuntimeBlockSequence.Add(1))
  477. return openAIAccountRuntimeBlockSnapshot{}
  478. }
  479. generation, _ := s.openaiAccountRuntimeBlockGeneration.Load(account.ID)
  480. gen, _ := generation.(uint64)
  481. return openAIAccountRuntimeBlockSnapshot{until: until, generation: gen, blocked: true}
  482. }
  483. // clearOpenAIAccountRuntimeBlockIfUnchanged deletes the in-process account block
  484. // only when generation and deadline are unchanged. A newer block installed after
  485. // peek must be kept even if its deadline happens to match.
  486. func (s *OpenAIGatewayService) clearOpenAIAccountRuntimeBlockIfUnchanged(accountID int64, snapshot openAIAccountRuntimeBlockSnapshot) {
  487. if s == nil || accountID <= 0 || !snapshot.blocked {
  488. return
  489. }
  490. mu := s.openAIAccountRuntimeBlockLock(accountID)
  491. mu.Lock()
  492. defer mu.Unlock()
  493. generation, ok := s.openaiAccountRuntimeBlockGeneration.Load(accountID)
  494. if !ok || generation != snapshot.generation {
  495. return
  496. }
  497. current, ok := s.openaiAccountRuntimeBlockUntil.Load(accountID)
  498. currentUntil, isTime := current.(time.Time)
  499. if !ok || !isTime || !currentUntil.Equal(snapshot.until) {
  500. return
  501. }
  502. s.openaiAccountRuntimeBlockUntil.Delete(accountID)
  503. s.openaiOAuth429RetryStartedAt.Delete(accountID)
  504. s.openaiAccountRuntimeBlockGeneration.Store(accountID, s.openaiAccountRuntimeBlockSequence.Add(1))
  505. }
  506. // isOpenAIAccountRequestRuntimeBlocked treats persisted cooldown fields on the
  507. // scheduling Account as source of truth. When TempUnschedulableUntil,
  508. // RateLimitResetAt, and OverloadUntil are all inactive, a stale local account
  509. // block is dropped with generation+deadline CAS. Model-scoped transient blocks
  510. // are left alone. This is fail-open if a DB write failed or the snapshot has
  511. // not caught up yet: empty cooldown fields drop the local account-level block.
  512. func (s *OpenAIGatewayService) isOpenAIAccountRequestRuntimeBlocked(account *Account, requestedModel string) bool {
  513. if s == nil {
  514. return false
  515. }
  516. snapshot := s.peekOpenAIAccountRuntimeBlock(account)
  517. if snapshot.blocked {
  518. if accountPersistedSchedulingCooldownActive(account) {
  519. return true
  520. }
  521. s.clearOpenAIAccountRuntimeBlockIfUnchanged(account.ID, snapshot)
  522. }
  523. return s.isOpenAIAccountModelRuntimeBlocked(account, requestedModel)
  524. }
  525. func (s *OpenAIGatewayService) recordOpenAIOAuth429() {
  526. if s == nil {
  527. return
  528. }
  529. now := time.Now()
  530. windowStart := s.openaiOAuth429WindowStartUnixNano.Load()
  531. if windowStart == 0 || now.Sub(time.Unix(0, windowStart)) >= openAIOAuth429StormWindow {
  532. if s.openaiOAuth429WindowStartUnixNano.CompareAndSwap(windowStart, now.UnixNano()) {
  533. s.openaiOAuth429WindowCount.Store(1)
  534. return
  535. }
  536. }
  537. s.openaiOAuth429WindowCount.Add(1)
  538. }
  539. func (s *OpenAIGatewayService) ShouldStopOpenAIOAuth429Failover(account *Account, statusCode int, failedSwitches int, state *OpenAIOAuth429FailoverState) bool {
  540. if failedSwitches < openAIOAuth429StormMaxAccountSwitches {
  541. return false
  542. }
  543. if state != nil && state.grokOAuth429FollowupPending {
  544. // The follow-up budget was armed by a Grok OAuth 429. Consume it on
  545. // any failing follow-up account, even if a mixed pool selected an API-key
  546. // account next.
  547. return true
  548. }
  549. if isGrokOAuthAccount(account) {
  550. if state == nil {
  551. // Preserve the old threshold for callers that have not adopted the
  552. // request-local state contract yet.
  553. return statusCode == http.StatusTooManyRequests && failedSwitches >= 2
  554. }
  555. if statusCode == http.StatusTooManyRequests {
  556. state.grokOAuth429FollowupPending = true
  557. }
  558. return false
  559. }
  560. if statusCode != http.StatusTooManyRequests || !isOpenAIOAuthAccount(account) {
  561. return false
  562. }
  563. // Each OpenAI OAuth candidate has already consumed its full same-account
  564. // retry window before reaching this switch point. A global storm is useful
  565. // telemetry, but must not prevent trying the bounded next-account budget.
  566. return failedSwitches >= openAIOAuth429MaxAccountAttempts
  567. }