openai_account_runtime_block_fastpath_test.go 38 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "errors"
  6. "net/http"
  7. "testing"
  8. "time"
  9. "github.com/Wei-Shaw/sub2api/internal/config"
  10. "github.com/stretchr/testify/require"
  11. )
  12. type oauth429RateLimitRepo struct {
  13. mockAccountRepoForGemini
  14. setRateLimitedCalls int
  15. lastRateLimitedUntil time.Time
  16. setModelRateLimitCalls int
  17. lastModelRateLimitKey string
  18. lastModelRateLimitedUntil time.Time
  19. }
  20. func (r *oauth429RateLimitRepo) SetRateLimited(_ context.Context, _ int64, until time.Time) error {
  21. r.setRateLimitedCalls++
  22. r.lastRateLimitedUntil = until
  23. return nil
  24. }
  25. func (r *oauth429RateLimitRepo) SetModelRateLimit(_ context.Context, _ int64, scope string, until time.Time, _ ...string) error {
  26. r.setModelRateLimitCalls++
  27. r.lastModelRateLimitKey = scope
  28. r.lastModelRateLimitedUntil = until
  29. return nil
  30. }
  31. func TestOpenAI429FastPath_KeepsOAuthAccountSchedulableDuringRetryWindow(t *testing.T) {
  32. repo := &oauth429RateLimitRepo{}
  33. rateLimits := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
  34. svc := &OpenAIGatewayService{rateLimitService: rateLimits}
  35. rateLimits.SetAccountRuntimeBlocker(svc)
  36. account := &Account{ID: 42, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  37. apiKeyAccount := &Account{ID: 43, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  38. setupTokenAccount := &Account{ID: 44, Platform: PlatformOpenAI, Type: AccountTypeSetupToken}
  39. grokOAuthAccount := &Account{ID: 45, Platform: PlatformGrok, Type: AccountTypeOAuth}
  40. shouldDisable := svc.handleOpenAIAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{}, nil)
  41. apiKeyShouldDisable := svc.handleOpenAIAccountUpstreamError(context.Background(), apiKeyAccount, http.StatusTooManyRequests, http.Header{}, nil)
  42. require.False(t, shouldDisable)
  43. require.False(t, apiKeyShouldDisable)
  44. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  45. require.True(t, svc.isOpenAIAccountRuntimeBlocked(apiKeyAccount), "API-key 429 keeps the existing scheduler cooldown behavior")
  46. require.Equal(t, 1, repo.setRateLimitedCalls, "only the API-key 429 should persist a scheduler block")
  47. require.True(t, svc.shouldRetryOpenAIOAuth429OnSameAccount(account, http.StatusTooManyRequests, false))
  48. require.True(t, svc.shouldRetryOpenAIOAuth429OnSameAccount(setupTokenAccount, http.StatusTooManyRequests, false))
  49. require.False(t, svc.shouldRetryOpenAIOAuth429OnSameAccount(apiKeyAccount, http.StatusTooManyRequests, false))
  50. require.False(t, svc.shouldRetryOpenAIOAuth429OnSameAccount(grokOAuthAccount, http.StatusTooManyRequests, false))
  51. require.WithinDuration(t, time.Now().Add(openAIOAuth429RetryWindow), svc.openAIOAuth429RetryDeadline(account), time.Second)
  52. require.WithinDuration(t, time.Now().Add(openAIOAuth429RetryWindow), svc.openAIOAuth429RetryDeadline(setupTokenAccount), time.Second)
  53. }
  54. func TestOpenAI429FastPath_BlocksOAuthOnlyAfterRetryWindow(t *testing.T) {
  55. svc := &OpenAIGatewayService{}
  56. account := &Account{ID: 420, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  57. svc.openaiOAuth429RetryStartedAt.Store(account.ID, time.Now().Add(-openAIOAuth429RetryWindow-time.Second))
  58. svc.markOpenAIOAuth429RateLimited(context.Background(), account, http.Header{}, nil)
  59. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  60. require.False(t, svc.shouldRetryOpenAIOAuth429OnSameAccount(account, http.StatusTooManyRequests, false))
  61. }
  62. func TestOpenAI429FastPath_DoesNotBlockOAuthWhenFallbackDisabled(t *testing.T) {
  63. repo := &oauth429RateLimitRepo{}
  64. settingRepo := newMockSettingRepo()
  65. settingRepo.data[SettingKeyRateLimit429CooldownSettings] = `{"enabled":false,"cooldown_seconds":12}`
  66. rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
  67. rateLimitService.SetSettingService(NewSettingService(settingRepo, &config.Config{}))
  68. svc := &OpenAIGatewayService{rateLimitService: rateLimitService}
  69. rateLimitService.SetAccountRuntimeBlocker(svc)
  70. account := &Account{ID: 425, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  71. svc.openaiOAuth429RetryStartedAt.Store(account.ID, time.Now().Add(-openAIOAuth429RetryWindow-time.Second))
  72. svc.markOpenAIOAuth429RateLimited(context.Background(), account, http.Header{}, []byte(`{"detail":"Rate limit exceeded"}`))
  73. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account), "disabled 429 fallback must not create an OAuth runtime cooldown")
  74. require.Zero(t, repo.setRateLimitedCalls, "disabled 429 fallback must not persist a scheduler cooldown")
  75. }
  76. func TestOpenAI429FastPath_DoesNotBlockOAuthWhenQuotaWindowIsNotExhausted(t *testing.T) {
  77. repo := &oauth429RateLimitRepo{}
  78. settingRepo := newMockSettingRepo()
  79. settingRepo.data[SettingKeyRateLimit429CooldownSettings] = `{"enabled":false,"cooldown_seconds":12}`
  80. rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
  81. rateLimitService.SetSettingService(NewSettingService(settingRepo, &config.Config{}))
  82. svc := &OpenAIGatewayService{rateLimitService: rateLimitService}
  83. rateLimitService.SetAccountRuntimeBlocker(svc)
  84. account := &Account{ID: 426, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  85. svc.openaiOAuth429RetryStartedAt.Store(account.ID, time.Now().Add(-openAIOAuth429RetryWindow-time.Second))
  86. headers := http.Header{}
  87. headers.Set("x-codex-primary-used-percent", "37")
  88. headers.Set("x-codex-primary-reset-after-seconds", "604800")
  89. headers.Set("x-codex-primary-window-minutes", "10080")
  90. headers.Set("x-codex-secondary-used-percent", "20")
  91. headers.Set("x-codex-secondary-reset-after-seconds", "3600")
  92. headers.Set("x-codex-secondary-window-minutes", "300")
  93. svc.handleOpenAIAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, headers, []byte(`{"detail":"Rate limit exceeded"}`))
  94. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account), "non-exhausted quota headers must use the configurable fallback")
  95. require.Zero(t, repo.setRateLimitedCalls, "disabled 429 fallback must not persist a scheduler cooldown")
  96. }
  97. func TestOpenAI429FastPath_BlocksOAuthImmediatelyWhenSevenDayQuotaIsExhausted(t *testing.T) {
  98. repo := &oauth429RateLimitRepo{}
  99. rateLimits := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
  100. svc := &OpenAIGatewayService{rateLimitService: rateLimits}
  101. rateLimits.SetAccountRuntimeBlocker(svc)
  102. account := &Account{ID: 423, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  103. headers := http.Header{}
  104. headers.Set("x-codex-primary-used-percent", "100")
  105. headers.Set("x-codex-primary-reset-after-seconds", "604800")
  106. headers.Set("x-codex-primary-window-minutes", "10080")
  107. headers.Set("x-codex-secondary-used-percent", "20")
  108. headers.Set("x-codex-secondary-reset-after-seconds", "3600")
  109. headers.Set("x-codex-secondary-window-minutes", "300")
  110. shouldDisable := svc.handleOpenAIAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, headers, []byte(`{"error":{"type":"rate_limit_error","code":"rate_limit_exceeded"}}`))
  111. require.False(t, shouldDisable)
  112. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  113. require.Equal(t, 1, repo.setRateLimitedCalls)
  114. require.Greater(t, time.Until(repo.lastRateLimitedUntil), 6*24*time.Hour)
  115. require.False(t, svc.ShouldRetryOpenAIOAuth429(account, headers, nil))
  116. }
  117. func TestOpenAI429FastPath_SparkQuotaOnlyBlocksSparkModel(t *testing.T) {
  118. repo := &oauth429RateLimitRepo{}
  119. rateLimits := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
  120. svc := &OpenAIGatewayService{rateLimitService: rateLimits}
  121. rateLimits.SetAccountRuntimeBlocker(svc)
  122. account := &Account{ID: 425, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  123. headers := http.Header{}
  124. headers.Set("x-codex-primary-used-percent", "100")
  125. headers.Set("x-codex-primary-reset-after-seconds", "604800")
  126. headers.Set("x-codex-primary-window-minutes", "10080")
  127. headers.Set("x-codex-secondary-used-percent", "20")
  128. headers.Set("x-codex-secondary-reset-after-seconds", "3600")
  129. headers.Set("x-codex-secondary-window-minutes", "300")
  130. shouldDisable := svc.handleOpenAIAccountUpstreamError(
  131. context.Background(), account, http.StatusTooManyRequests, headers,
  132. []byte(`{"error":{"type":"rate_limit_error","code":"rate_limit_exceeded"}}`),
  133. "gpt-5.3-codex-spark",
  134. )
  135. require.False(t, shouldDisable)
  136. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account), "Spark quota must not create an account runtime block")
  137. require.Equal(t, 0, repo.setRateLimitedCalls, "Spark quota must not persist account-level rate limit")
  138. require.Equal(t, 1, repo.setModelRateLimitCalls)
  139. require.Equal(t, "gpt-5.3-codex-spark", repo.lastModelRateLimitKey)
  140. require.Greater(t, time.Until(repo.lastModelRateLimitedUntil), 6*24*time.Hour)
  141. }
  142. func TestOpenAI429FastPath_SparkTransient429UsesShortFallback(t *testing.T) {
  143. repo := &oauth429RateLimitRepo{}
  144. rateLimits := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
  145. svc := &OpenAIGatewayService{rateLimitService: rateLimits}
  146. rateLimits.SetAccountRuntimeBlocker(svc)
  147. account := &Account{ID: 428, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  148. headers := http.Header{}
  149. headers.Set("x-codex-primary-used-percent", "37")
  150. headers.Set("x-codex-primary-reset-after-seconds", "604800")
  151. headers.Set("x-codex-primary-window-minutes", "10080")
  152. headers.Set("x-codex-secondary-used-percent", "20")
  153. headers.Set("x-codex-secondary-reset-after-seconds", "3600")
  154. headers.Set("x-codex-secondary-window-minutes", "300")
  155. shouldDisable := svc.handleOpenAIAccountUpstreamError(
  156. context.Background(), account, http.StatusTooManyRequests, headers,
  157. []byte(`{"error":{"type":"rate_limit_error","code":"rate_limit_exceeded"}}`),
  158. "gpt-5.3-codex-spark",
  159. )
  160. require.False(t, shouldDisable)
  161. require.Equal(t, 1, repo.setModelRateLimitCalls)
  162. require.Less(t, time.Until(repo.lastModelRateLimitedUntil), time.Minute)
  163. require.Greater(t, time.Until(repo.lastModelRateLimitedUntil), time.Second)
  164. }
  165. func TestOpenAIStream429_SparkQuotaUsesQuotaHeaders(t *testing.T) {
  166. repo := &oauth429RateLimitRepo{}
  167. rateLimits := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
  168. svc := &OpenAIGatewayService{rateLimitService: rateLimits}
  169. rateLimits.SetAccountRuntimeBlocker(svc)
  170. account := &Account{ID: 429, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  171. headers := http.Header{}
  172. headers.Set("x-codex-primary-used-percent", "100")
  173. headers.Set("x-codex-primary-reset-after-seconds", "604800")
  174. headers.Set("x-codex-primary-window-minutes", "10080")
  175. headers.Set("x-codex-secondary-used-percent", "20")
  176. headers.Set("x-codex-secondary-reset-after-seconds", "3600")
  177. headers.Set("x-codex-secondary-window-minutes", "300")
  178. payload := []byte(`{"type":"error","error":{"type":"rate_limit_error","code":"rate_limit_exceeded"}}`)
  179. status, shouldDisable := svc.handleOpenAIStreamTerminalAccountSideEffects(
  180. nil, account, payload, "quota exhausted", headers, "gpt-5.3-codex-spark",
  181. )
  182. require.Equal(t, http.StatusTooManyRequests, status)
  183. require.False(t, shouldDisable)
  184. require.Equal(t, 1, repo.setModelRateLimitCalls)
  185. require.Equal(t, "gpt-5.3-codex-spark", repo.lastModelRateLimitKey)
  186. require.Greater(t, time.Until(repo.lastModelRateLimitedUntil), 6*24*time.Hour)
  187. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  188. }
  189. func TestOpenAIStreamFailover_Spark429KeepsModelScope(t *testing.T) {
  190. repo := &oauth429RateLimitRepo{}
  191. rateLimits := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
  192. svc := &OpenAIGatewayService{rateLimitService: rateLimits}
  193. rateLimits.SetAccountRuntimeBlocker(svc)
  194. account := &Account{ID: 432, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  195. headers := http.Header{}
  196. headers.Set("x-codex-primary-used-percent", "100")
  197. headers.Set("x-codex-primary-reset-after-seconds", "604800")
  198. headers.Set("x-codex-primary-window-minutes", "10080")
  199. headers.Set("x-codex-secondary-used-percent", "20")
  200. headers.Set("x-codex-secondary-reset-after-seconds", "3600")
  201. headers.Set("x-codex-secondary-window-minutes", "300")
  202. payload := []byte(`{"type":"error","error":{"type":"rate_limit_error","code":"rate_limit_exceeded"}}`)
  203. failoverErr := svc.newOpenAIStreamFailoverErrorWithModel(
  204. nil, account, false, "", payload, "quota exhausted", "gpt-5.3-codex-spark", headers,
  205. )
  206. require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
  207. require.Equal(t, 0, repo.setRateLimitedCalls)
  208. require.Equal(t, 1, repo.setModelRateLimitCalls)
  209. require.Equal(t, "gpt-5.3-codex-spark", repo.lastModelRateLimitKey)
  210. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  211. }
  212. func TestOpenAIWSErrorEvent_OrdinaryModelIgnoresHandshakeQuotaHeaders(t *testing.T) {
  213. repo := &oauth429RateLimitRepo{}
  214. rateLimits := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
  215. svc := &OpenAIGatewayService{rateLimitService: rateLimits}
  216. rateLimits.SetAccountRuntimeBlocker(svc)
  217. account := &Account{ID: 430, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  218. headers := http.Header{}
  219. headers.Set("x-codex-primary-used-percent", "100")
  220. headers.Set("x-codex-primary-reset-after-seconds", "604800")
  221. headers.Set("x-codex-primary-window-minutes", "10080")
  222. payload := []byte(`{"type":"error","error":{"type":"rate_limit_error","code":"rate_limit_exceeded"}}`)
  223. svc.persistOpenAIWSRateLimitSignal(context.Background(), account, headers, payload, "rate_limit_exceeded", "rate_limit_error", "quota exhausted", "gpt-5.3-codex")
  224. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  225. require.Zero(t, repo.setRateLimitedCalls)
  226. }
  227. func TestOpenAIWSErrorEvent_SparkQuotaUsesHandshakeQuotaHeaders(t *testing.T) {
  228. repo := &oauth429RateLimitRepo{}
  229. rateLimits := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
  230. svc := &OpenAIGatewayService{rateLimitService: rateLimits}
  231. rateLimits.SetAccountRuntimeBlocker(svc)
  232. account := &Account{ID: 431, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  233. headers := http.Header{}
  234. headers.Set("x-codex-primary-used-percent", "100")
  235. headers.Set("x-codex-primary-reset-after-seconds", "604800")
  236. headers.Set("x-codex-primary-window-minutes", "10080")
  237. payload := []byte(`{"type":"error","error":{"type":"rate_limit_error","code":"rate_limit_exceeded"}}`)
  238. svc.persistOpenAIWSRateLimitSignal(context.Background(), account, headers, payload, "rate_limit_exceeded", "rate_limit_error", "quota exhausted", "gpt-5.3-codex-spark")
  239. require.Equal(t, 1, repo.setModelRateLimitCalls)
  240. require.Equal(t, "gpt-5.3-codex-spark", repo.lastModelRateLimitKey)
  241. require.Greater(t, time.Until(repo.lastModelRateLimitedUntil), 6*24*time.Hour)
  242. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  243. }
  244. func TestOpenAI429FastPath_SparkShadowQuotaStaysModelScoped(t *testing.T) {
  245. repo := &oauth429RateLimitRepo{}
  246. rateLimits := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
  247. svc := &OpenAIGatewayService{rateLimitService: rateLimits}
  248. rateLimits.SetAccountRuntimeBlocker(svc)
  249. parentID := int64(426)
  250. shadow := &Account{ID: 427, Platform: PlatformOpenAI, Type: AccountTypeOAuth, ParentAccountID: &parentID, QuotaDimension: QuotaDimensionSpark}
  251. headers := http.Header{}
  252. headers.Set("x-codex-primary-used-percent", "100")
  253. headers.Set("x-codex-primary-reset-after-seconds", "604800")
  254. headers.Set("x-codex-primary-window-minutes", "10080")
  255. shouldDisable := svc.handleOpenAIAccountUpstreamError(
  256. context.Background(), shadow, http.StatusTooManyRequests, headers,
  257. []byte(`{"error":{"type":"rate_limit_error","code":"rate_limit_exceeded"}}`),
  258. "gpt-5.3-codex-spark",
  259. )
  260. require.False(t, shouldDisable)
  261. require.False(t, svc.isOpenAIAccountRuntimeBlocked(shadow))
  262. require.Equal(t, 0, repo.setRateLimitedCalls)
  263. require.Equal(t, 1, repo.setModelRateLimitCalls)
  264. require.Equal(t, "gpt-5.3-codex-spark", repo.lastModelRateLimitKey)
  265. }
  266. func TestOpenAI429FastPath_RetriesOAuthWhenNoQuotaSignalExists(t *testing.T) {
  267. svc := &OpenAIGatewayService{}
  268. account := &Account{ID: 424, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  269. headers := http.Header{}
  270. require.True(t, svc.ShouldRetryOpenAIOAuth429(account, headers, []byte(`{"error":{"type":"rate_limit_error","message":"try again"}}`)))
  271. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  272. }
  273. func TestOpenAI429FastPath_ExplicitRetryAfterWaitsBeforeScheduling(t *testing.T) {
  274. repo := &oauth429RateLimitRepo{}
  275. limits := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
  276. svc := &OpenAIGatewayService{rateLimitService: limits}
  277. limits.SetAccountRuntimeBlocker(svc)
  278. a := &Account{ID: 425, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  279. headers := http.Header{"Retry-After": []string{"90"}}
  280. body := []byte(`{"error":{"type":"rate_limit_error","message":"slow down"}}`)
  281. require.False(t, svc.ShouldRetryOpenAIOAuth429(a, headers, body))
  282. svc.handleOpenAIAccountUpstreamError(context.Background(), a, 429, headers, body)
  283. require.True(t, svc.isOpenAIAccountRuntimeBlocked(a))
  284. require.Equal(t, 1, repo.setRateLimitedCalls)
  285. until, ok := svc.openaiAccountRuntimeBlockUntil.Load(a.ID)
  286. require.True(t, ok)
  287. require.Greater(t, time.Until(until.(time.Time)), 89*time.Second)
  288. }
  289. func TestOpenAIStream429IgnoresSuccessfulQuotaSnapshotHeaders(t *testing.T) {
  290. repo := &oauth429RateLimitRepo{}
  291. rateLimits := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
  292. svc := &OpenAIGatewayService{rateLimitService: rateLimits}
  293. rateLimits.SetAccountRuntimeBlocker(svc)
  294. account := &Account{ID: 421, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  295. svc.openaiOAuth429RetryStartedAt.Store(account.ID, time.Now().Add(-openAIOAuth429RetryWindow-time.Second))
  296. headers := http.Header{}
  297. headers.Set("x-codex-primary-used-percent", "37")
  298. headers.Set("x-codex-primary-reset-after-seconds", "604800")
  299. headers.Set("x-codex-primary-window-minutes", "10080")
  300. payload := []byte(`{"type":"error","error":{"type":"rate_limit_error","code":"rate_limit_exceeded","message":"slow down"}}`)
  301. status, disabled := svc.handleOpenAIStreamTerminalAccountSideEffects(nil, account, payload, "slow down", headers)
  302. require.Equal(t, http.StatusTooManyRequests, status)
  303. require.False(t, disabled)
  304. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  305. value, ok := svc.openaiAccountRuntimeBlockUntil.Load(account.ID)
  306. require.True(t, ok)
  307. blockedUntil, ok := value.(time.Time)
  308. require.True(t, ok)
  309. require.Less(t, time.Until(blockedUntil), time.Minute, "stream 429 must not inherit the normal seven-day quota snapshot")
  310. if !repo.lastRateLimitedUntil.IsZero() {
  311. require.Less(t, time.Until(repo.lastRateLimitedUntil), time.Minute)
  312. }
  313. }
  314. func TestOpenAIHTTP429StillUsesQuotaResetHeaders(t *testing.T) {
  315. svc := &OpenAIGatewayService{rateLimitService: &RateLimitService{}}
  316. account := &Account{ID: 422, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  317. svc.openaiOAuth429RetryStartedAt.Store(account.ID, time.Now().Add(-openAIOAuth429RetryWindow-time.Second))
  318. headers := http.Header{}
  319. headers.Set("x-codex-primary-used-percent", "100")
  320. headers.Set("x-codex-primary-reset-after-seconds", "604800")
  321. headers.Set("x-codex-primary-window-minutes", "10080")
  322. svc.markOpenAIOAuth429RateLimited(context.Background(), account, headers, nil)
  323. value, ok := svc.openaiAccountRuntimeBlockUntil.Load(account.ID)
  324. require.True(t, ok)
  325. blockedUntil, ok := value.(time.Time)
  326. require.True(t, ok)
  327. require.Greater(t, time.Until(blockedUntil), 6*24*time.Hour, "real HTTP 429 must retain the upstream quota reset")
  328. }
  329. func TestOpenAI429RetryDelayHonorsFullRetryAfter(t *testing.T) {
  330. deadline := time.Now().Add(openAIOAuth429RetryWindow)
  331. require.Equal(t, openAIOAuth429RetryDelay, openAIOAuth429SameAccountRetryDelay(nil, deadline))
  332. require.Equal(t, 90*time.Second, openAIOAuth429SameAccountRetryDelay(http.Header{"Retry-After": []string{"90"}}, deadline))
  333. require.Equal(t, 90*time.Second, openAIOAuth429SameAccountRetryDelay(http.Header{"Retry-After": []string{"90"}}, time.Now().Add(time.Second)))
  334. }
  335. func TestOpenAI429FastPath_OpenCodeGoUsageLimitUsesMessageResetDuration(t *testing.T) {
  336. repo := &rateLimit429AccountRepoStub{}
  337. rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
  338. svc := &OpenAIGatewayService{rateLimitService: rateLimitService}
  339. rateLimitService.SetAccountRuntimeBlocker(svc)
  340. account := &Account{ID: 44, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  341. body := []byte(`{"type":"error","error":{"type":"GoUsageLimitError","message":"5-hour usage limit reached. Resets in 4hr 59min. To continue using this model now, enable usage from your available balance: https://opencode.ai/workspace/wrk_test/go"},"metadata":{"workspace":"wrk_test","limitName":"5 hour"}}`)
  342. before := time.Now()
  343. shouldDisable := svc.handleOpenAIAccountUpstreamError(
  344. context.Background(),
  345. account,
  346. http.StatusTooManyRequests,
  347. http.Header{},
  348. body,
  349. )
  350. after := time.Now()
  351. require.False(t, shouldDisable)
  352. require.Equal(t, 1, repo.rateLimitCalls)
  353. require.Equal(t, account.ID, repo.lastRateLimitID)
  354. expectedResetAfter := 4*time.Hour + 59*time.Minute
  355. require.False(t, repo.lastRateLimitReset.Before(before.Add(expectedResetAfter-time.Second)))
  356. require.False(t, repo.lastRateLimitReset.After(after.Add(expectedResetAfter)))
  357. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  358. }
  359. // TestOpenAI429FastPath_SkipsSparkShadow 外审第8轮 P1:spark 影子被选中后若 /responses 返回 429,
  360. // 不得按 global x-codex-* 信号写内存运行时熔断(否则 spark 被冷却到 global reset、单影子场景无可用账号)。
  361. func TestOpenAI429FastPath_SkipsSparkShadow(t *testing.T) {
  362. svc := &OpenAIGatewayService{}
  363. parentID := int64(800)
  364. shadow := &Account{
  365. ID: 801,
  366. Platform: PlatformOpenAI,
  367. Type: AccountTypeOAuth,
  368. ParentAccountID: &parentID,
  369. QuotaDimension: QuotaDimensionSpark,
  370. }
  371. normal := &Account{ID: 802, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  372. headers := http.Header{}
  373. headers.Set("x-codex-primary-used-percent", "100")
  374. headers.Set("x-codex-primary-reset-after-seconds", "18000")
  375. headers.Set("x-codex-primary-window-minutes", "300")
  376. svc.markOpenAIOAuth429RateLimited(context.Background(), shadow, headers, nil)
  377. svc.markOpenAIOAuth429RateLimited(context.Background(), normal, headers, nil)
  378. require.False(t, svc.isOpenAIAccountRuntimeBlocked(shadow), "spark shadow must not be runtime-blocked by /responses global 429")
  379. require.True(t, svc.isOpenAIAccountRuntimeBlocked(normal), "normal OpenAI OAuth account with an exhausted 5h window must be paused")
  380. }
  381. func TestOpenAIRuntimeBlock_AppliesToOpenAIAPIKeyWhenRateLimitServiceStopsScheduling(t *testing.T) {
  382. svc := &OpenAIGatewayService{}
  383. account := &Account{ID: 44, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  384. svc.BlockAccountScheduling(account, time.Time{}, "custom_error_code")
  385. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  386. }
  387. func TestOpenAIRuntimeBlock_DoesNotApplyToOtherPlatforms(t *testing.T) {
  388. svc := &OpenAIGatewayService{}
  389. account := &Account{ID: 45, Platform: PlatformGemini, Type: AccountTypeOAuth}
  390. svc.BlockAccountScheduling(account, time.Time{}, "custom_error_code")
  391. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  392. }
  393. func TestOpenAIRuntimeBlocker_IgnoresNonOpenAIFromRateLimitService(t *testing.T) {
  394. gateway := &OpenAIGatewayService{}
  395. repo := &rateLimitAccountRepoStub{}
  396. rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
  397. rateLimitService.SetAccountRuntimeBlocker(gateway)
  398. account := &Account{ID: 45, Platform: PlatformGemini, Type: AccountTypeOAuth}
  399. shouldDisable := rateLimitService.HandleUpstreamError(context.Background(), account, http.StatusForbidden, http.Header{}, []byte("forbidden"))
  400. require.True(t, shouldDisable)
  401. require.False(t, gateway.isOpenAIAccountRuntimeBlocked(account))
  402. }
  403. // 自 #4547(issue 4527 第4点)起,临时不可调度规则命中已知模型时按模型隔离:
  404. // 只封 (账号, 模型) 对,不再账号级一刀切;未知模型仍走账号级兜底
  405. // (见 TestOpenAITempUnschedulable_UnknownModelKeepsAccountRuntimeBlock)。
  406. // 池模式规则仍然生效(issue 4470):停止同账号重试并对命中模型设临时封锁。
  407. func TestOpenAIPoolModeTempRule_StopsSameAccountRetryAndIsolatesBlockToModel(t *testing.T) {
  408. repo := &errorPolicyRepoStub{}
  409. rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
  410. gateway := &OpenAIGatewayService{
  411. cfg: &config.Config{},
  412. rateLimitService: rateLimitService,
  413. }
  414. rateLimitService.SetAccountRuntimeBlocker(gateway)
  415. account := &Account{
  416. ID: 46,
  417. Platform: PlatformOpenAI,
  418. Type: AccountTypeAPIKey,
  419. Status: StatusActive,
  420. Schedulable: true,
  421. Credentials: map[string]any{
  422. "pool_mode": true,
  423. "pool_mode_retry_status_codes": []any{float64(http.StatusServiceUnavailable)},
  424. "temp_unschedulable_enabled": true,
  425. "temp_unschedulable_rules": []any{
  426. map[string]any{
  427. "error_code": float64(http.StatusServiceUnavailable),
  428. "keywords": []any{"unavailable"},
  429. "duration_minutes": float64(30),
  430. },
  431. },
  432. },
  433. }
  434. body := []byte(`{"error":{"message":"Service temporarily unavailable"}}`)
  435. resp := &http.Response{
  436. StatusCode: http.StatusServiceUnavailable,
  437. Header: http.Header{},
  438. }
  439. failoverErr := gateway.failoverOpenAIUpstreamHTTPError(
  440. context.Background(),
  441. nil,
  442. account,
  443. resp,
  444. body,
  445. "Service temporarily unavailable",
  446. "gpt-5.4",
  447. )
  448. require.NotNil(t, failoverErr)
  449. require.False(t, failoverErr.RetryableOnSameAccount)
  450. require.Zero(t, repo.tempCalls)
  451. require.Equal(t, 0, repo.setErrCalls)
  452. require.Equal(t, StatusActive, account.Status)
  453. require.Len(t, repo.modelRateLimitCalls, 1)
  454. require.Equal(t, "gpt-5.4", repo.modelRateLimitCalls[0].scope)
  455. require.False(t, gateway.isOpenAIAccountRuntimeBlocked(account))
  456. require.False(t, gateway.isOpenAIAccountRequestRuntimeBlocked(account, "gpt-5.5"))
  457. }
  458. func TestOpenAIPoolModeRetryable5xx_DoesNotCreateModelTransientBlock(t *testing.T) {
  459. repo := &errorPolicyRepoStub{}
  460. rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
  461. gateway := &OpenAIGatewayService{rateLimitService: rateLimitService}
  462. account := &Account{
  463. ID: 47,
  464. Platform: PlatformOpenAI,
  465. Type: AccountTypeAPIKey,
  466. Credentials: map[string]any{
  467. "pool_mode": true,
  468. "pool_mode_retry_status_codes": []any{float64(524)},
  469. },
  470. }
  471. for i := 0; i < 2; i++ {
  472. shouldDisable := gateway.handleOpenAIAccountUpstreamError(
  473. context.Background(),
  474. account,
  475. 524,
  476. http.Header{},
  477. []byte(`{"error":{"message":"upstream timeout"}}`),
  478. "gpt-5.4",
  479. )
  480. require.False(t, shouldDisable)
  481. }
  482. require.False(t, gateway.isOpenAIAccountRequestRuntimeBlocked(account, "gpt-5.4"))
  483. }
  484. func TestOpenAIPoolModeNonRetryable5xx_StillCreatesModelTransientBlock(t *testing.T) {
  485. repo := &errorPolicyRepoStub{}
  486. rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
  487. gateway := &OpenAIGatewayService{rateLimitService: rateLimitService}
  488. account := &Account{
  489. ID: 48,
  490. Platform: PlatformOpenAI,
  491. Type: AccountTypeAPIKey,
  492. Credentials: map[string]any{
  493. "pool_mode": true,
  494. "pool_mode_retry_status_codes": []any{float64(http.StatusGatewayTimeout)},
  495. },
  496. }
  497. for i := 0; i < 2; i++ {
  498. shouldDisable := gateway.handleOpenAIAccountUpstreamError(
  499. context.Background(),
  500. account,
  501. http.StatusServiceUnavailable,
  502. http.Header{},
  503. []byte(`{"error":{"message":"upstream unavailable"}}`),
  504. "gpt-5.4",
  505. )
  506. require.False(t, shouldDisable)
  507. }
  508. require.True(t, gateway.isOpenAIAccountRequestRuntimeBlocked(account, "gpt-5.4"))
  509. }
  510. func TestOpenAINonPoolAPIKey5xx_StillCreatesModelTransientBlock(t *testing.T) {
  511. repo := &errorPolicyRepoStub{}
  512. rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
  513. gateway := &OpenAIGatewayService{rateLimitService: rateLimitService}
  514. account := &Account{
  515. ID: 49,
  516. Platform: PlatformOpenAI,
  517. Type: AccountTypeAPIKey,
  518. }
  519. for i := 0; i < 2; i++ {
  520. shouldDisable := gateway.handleOpenAIAccountUpstreamError(
  521. context.Background(),
  522. account,
  523. http.StatusGatewayTimeout,
  524. http.Header{},
  525. []byte(`{"error":{"message":"upstream timeout"}}`),
  526. "gpt-5.4",
  527. )
  528. require.False(t, shouldDisable)
  529. }
  530. require.True(t, gateway.isOpenAIAccountRequestRuntimeBlocked(account, "gpt-5.4"))
  531. }
  532. func TestOpenAIModelNotFound_DoesNotRuntimeBlockWholeAccount(t *testing.T) {
  533. repo := &modelNotFoundAccountRepoStub{}
  534. svc := &OpenAIGatewayService{
  535. rateLimitService: &RateLimitService{accountRepo: repo},
  536. }
  537. account := openAIModelNotFoundTempAccount()
  538. shouldDisable := svc.handleOpenAIAccountUpstreamError(
  539. context.Background(),
  540. account,
  541. http.StatusNotFound,
  542. http.Header{},
  543. []byte(`{"error":{"code":"model_not_found","message":"model not found"}}`),
  544. "gpt-5.4",
  545. )
  546. require.True(t, shouldDisable)
  547. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  548. require.Zero(t, repo.tempCalls)
  549. require.Len(t, repo.modelRateLimitCalls, 1)
  550. }
  551. func TestOpenAIModelTempUnschedulable_DoesNotRuntimeBlockWholeAccount(t *testing.T) {
  552. repo := &modelNotFoundAccountRepoStub{}
  553. svc := &OpenAIGatewayService{
  554. rateLimitService: &RateLimitService{accountRepo: repo},
  555. }
  556. account := openAIModelNotFoundTempAccount()
  557. shouldDisable := svc.handleOpenAIAccountUpstreamError(
  558. context.Background(),
  559. account,
  560. http.StatusNotFound,
  561. http.Header{},
  562. []byte(`{"error":{"message":"endpoint not found"}}`),
  563. "gpt-5.4",
  564. )
  565. require.True(t, shouldDisable)
  566. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  567. require.Zero(t, repo.tempCalls)
  568. require.Len(t, repo.modelRateLimitCalls, 1)
  569. require.Equal(t, "gpt-5.4", repo.modelRateLimitCalls[0].scope)
  570. }
  571. func TestOpenAIModelTempUnschedulable_WriteFailureDoesNotRuntimeBlockWholeAccount(t *testing.T) {
  572. repo := &modelNotFoundAccountRepoStub{modelRateLimitErr: errors.New("write failed")}
  573. svc := &OpenAIGatewayService{
  574. rateLimitService: &RateLimitService{accountRepo: repo},
  575. }
  576. account := openAIModelNotFoundTempAccount()
  577. shouldDisable := svc.handleOpenAIAccountUpstreamError(
  578. context.Background(),
  579. account,
  580. http.StatusNotFound,
  581. http.Header{},
  582. []byte(`{"error":{"message":"endpoint not found"}}`),
  583. "gpt-5.4",
  584. )
  585. require.True(t, shouldDisable)
  586. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  587. require.Zero(t, repo.tempCalls)
  588. require.Len(t, repo.modelRateLimitCalls, 1)
  589. }
  590. func TestOpenAIOAuth429_MatchingModelTempRuleAvoidsAccountRuntimeBlock(t *testing.T) {
  591. repo := &modelNotFoundAccountRepoStub{}
  592. svc := &OpenAIGatewayService{
  593. rateLimitService: &RateLimitService{accountRepo: repo},
  594. }
  595. account := openAIModelNotFoundTempAccount()
  596. account.Type = AccountTypeOAuth
  597. account.Credentials["temp_unschedulable_rules"] = []any{
  598. map[string]any{
  599. "error_code": float64(http.StatusTooManyRequests),
  600. "keywords": []any{"model quota"},
  601. "duration_minutes": float64(10),
  602. },
  603. }
  604. shouldDisable := svc.handleOpenAIAccountUpstreamError(
  605. context.Background(),
  606. account,
  607. http.StatusTooManyRequests,
  608. http.Header{},
  609. []byte(`{"error":{"message":"model quota exhausted"}}`),
  610. "gpt-5.4",
  611. )
  612. require.True(t, shouldDisable)
  613. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  614. require.Len(t, repo.modelRateLimitCalls, 1)
  615. require.Equal(t, "gpt-5.4", repo.modelRateLimitCalls[0].scope)
  616. }
  617. func TestOpenAIOAuth429_NonmatchingModelTempRuleKeepsAccountRuntimeBlock(t *testing.T) {
  618. repo := &modelNotFoundAccountRepoStub{}
  619. svc := &OpenAIGatewayService{
  620. rateLimitService: &RateLimitService{accountRepo: repo},
  621. }
  622. account := openAIModelNotFoundTempAccount()
  623. account.Type = AccountTypeOAuth
  624. account.Credentials["temp_unschedulable_rules"] = []any{
  625. map[string]any{
  626. "error_code": float64(http.StatusTooManyRequests),
  627. "keywords": []any{"different marker"},
  628. "duration_minutes": float64(10),
  629. },
  630. }
  631. shouldDisable := svc.handleOpenAIAccountUpstreamError(
  632. context.Background(),
  633. account,
  634. http.StatusTooManyRequests,
  635. http.Header{},
  636. []byte(`{"error":{"message":"global rate limit"}}`),
  637. "gpt-5.4",
  638. )
  639. require.False(t, shouldDisable)
  640. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  641. require.True(t, svc.shouldRetryOpenAIOAuth429OnSameAccount(account, http.StatusTooManyRequests, false))
  642. require.Empty(t, repo.modelRateLimitCalls)
  643. }
  644. func TestOpenAITempUnschedulable_UnknownModelKeepsAccountRuntimeBlock(t *testing.T) {
  645. repo := &modelNotFoundAccountRepoStub{}
  646. svc := &OpenAIGatewayService{
  647. rateLimitService: &RateLimitService{accountRepo: repo},
  648. }
  649. account := openAIModelNotFoundTempAccount()
  650. shouldDisable := svc.handleOpenAIAccountUpstreamError(
  651. context.Background(),
  652. account,
  653. http.StatusNotFound,
  654. http.Header{},
  655. []byte(`{"error":{"message":"endpoint not found"}}`),
  656. )
  657. require.True(t, shouldDisable)
  658. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  659. require.Equal(t, 1, repo.tempCalls)
  660. require.Empty(t, repo.modelRateLimitCalls)
  661. }
  662. func TestOpenAIRuntimeBlock_DoesNotShortenExistingBlock(t *testing.T) {
  663. svc := &OpenAIGatewayService{}
  664. account := &Account{ID: 46, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  665. longUntil := time.Now().Add(10 * time.Minute)
  666. svc.BlockAccountScheduling(account, longUntil, "oauth_401")
  667. svc.BlockAccountScheduling(account, time.Time{}, "upstream_disable")
  668. value, ok := svc.openaiAccountRuntimeBlockUntil.Load(account.ID)
  669. require.True(t, ok)
  670. actualUntil, ok := value.(time.Time)
  671. require.True(t, ok)
  672. require.WithinDuration(t, longUntil, actualUntil, time.Second)
  673. }
  674. func TestOpenAIRuntimeBlock_ClearAccountSchedulingBlock(t *testing.T) {
  675. svc := &OpenAIGatewayService{}
  676. account := &Account{ID: 47, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  677. svc.BlockAccountScheduling(account, time.Now().Add(time.Minute), "429")
  678. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  679. svc.ClearAccountSchedulingBlock(account.ID)
  680. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  681. }
  682. func TestRuntimeBlockHonorsClearedPersistedCooldown(t *testing.T) {
  683. svc := &OpenAIGatewayService{}
  684. account := &Account{ID: 92, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true}
  685. svc.BlockAccountScheduling(account, time.Now().Add(30*time.Minute), "grok payment required")
  686. require.False(t, svc.isOpenAIAccountRequestRuntimeBlocked(account, "grok-3"))
  687. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  688. }
  689. func TestRuntimeBlockConditionalClearSkipsNewerGeneration(t *testing.T) {
  690. svc := &OpenAIGatewayService{}
  691. account := &Account{ID: 94, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true}
  692. firstUntil := time.Now().Add(10 * time.Minute)
  693. svc.BlockAccountScheduling(account, firstUntil, "stale")
  694. snapshot := svc.peekOpenAIAccountRuntimeBlock(account)
  695. require.True(t, snapshot.blocked)
  696. newerUntil := time.Now().Add(30 * time.Minute)
  697. svc.BlockAccountScheduling(account, newerUntil, "fresh")
  698. svc.clearOpenAIAccountRuntimeBlockIfUnchanged(account.ID, snapshot)
  699. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  700. require.False(t, svc.isOpenAIAccountRequestRuntimeBlocked(account, "grok-3"))
  701. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  702. }
  703. func TestRuntimeBlockKeepsActivePersistedCooldown(t *testing.T) {
  704. svc := &OpenAIGatewayService{}
  705. until := time.Now().Add(30 * time.Minute)
  706. account := &Account{
  707. ID: 93,
  708. Platform: PlatformGrok,
  709. Type: AccountTypeOAuth,
  710. Status: StatusActive,
  711. Schedulable: true,
  712. TempUnschedulableUntil: &until,
  713. }
  714. svc.BlockAccountScheduling(account, until, "grok payment required")
  715. require.True(t, svc.isOpenAIAccountRequestRuntimeBlocked(account, "grok-3"))
  716. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  717. }
  718. func TestShouldStopOpenAIOAuth429Failover_AfterBoundedFullWindows(t *testing.T) {
  719. svc := &OpenAIGatewayService{}
  720. account := &Account{ID: 42, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  721. apiKeyAccount := &Account{ID: 43, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  722. var state OpenAIOAuth429FailoverState
  723. require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1, &state))
  724. require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1, &state))
  725. require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 2, &state))
  726. require.True(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 3, &state))
  727. require.False(t, svc.ShouldStopOpenAIOAuth429Failover(apiKeyAccount, http.StatusTooManyRequests, 1, &state))
  728. require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusInternalServerError, 1, &state))
  729. require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 0, &state))
  730. }
  731. func TestShouldStopOpenAIOAuth429Failover_TracksOneGrokFollowupAttempt(t *testing.T) {
  732. svc := &OpenAIGatewayService{}
  733. account := &Account{ID: 44, Platform: PlatformGrok, Type: AccountTypeOAuth}
  734. apiKeyAccount := &Account{ID: 45, Platform: PlatformGrok, Type: AccountTypeAPIKey}
  735. t.Run("429 then 500 stops after one followup", func(t *testing.T) {
  736. var state OpenAIOAuth429FailoverState
  737. require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1, &state))
  738. require.True(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusInternalServerError, 2, &state))
  739. })
  740. t.Run("500 then 429 still allows one followup", func(t *testing.T) {
  741. var state OpenAIOAuth429FailoverState
  742. require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusInternalServerError, 1, &state))
  743. require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 2, &state))
  744. require.True(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusBadGateway, 3, &state))
  745. })
  746. t.Run("OAuth 429 then API-key failure consumes the same followup", func(t *testing.T) {
  747. var state OpenAIOAuth429FailoverState
  748. require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1, &state))
  749. require.True(t, svc.ShouldStopOpenAIOAuth429Failover(apiKeyAccount, http.StatusInternalServerError, 2, &state))
  750. })
  751. var state OpenAIOAuth429FailoverState
  752. require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 0, &state))
  753. require.False(t, svc.ShouldStopOpenAIOAuth429Failover(apiKeyAccount, http.StatusTooManyRequests, 2, &state))
  754. }