grok_credential_failure_test.go 67 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "errors"
  6. "net/http"
  7. "net/http/httptest"
  8. "sync"
  9. "testing"
  10. "time"
  11. infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
  12. "github.com/Wei-Shaw/sub2api/internal/pkg/xai"
  13. "github.com/gin-gonic/gin"
  14. "github.com/stretchr/testify/require"
  15. )
  16. type grokCredentialPersistingRepo struct {
  17. *tokenRefreshAccountRepo
  18. }
  19. func TestClassifyGrokCredentialFailureBillingExhaustionIsTransient(t *testing.T) {
  20. account := expiredGrokOAuthAccountForCredentialTest(9901)
  21. for _, message := range []string{
  22. "Grok OAuth refresh failed: spending limit reached",
  23. "included free usage exhausted",
  24. "credits exhausted",
  25. } {
  26. class := classifyGrokCredentialFailure(account, errors.New(message))
  27. require.Equal(t, GrokCredentialReasonRefreshTransient, class.reason, message)
  28. require.True(t, class.transient, message)
  29. require.False(t, class.permanent, message)
  30. }
  31. }
  32. func (r *grokCredentialPersistingRepo) SetError(ctx context.Context, id int64, message string) error {
  33. if err := ctx.Err(); err != nil {
  34. return err
  35. }
  36. if err := r.tokenRefreshAccountRepo.SetError(ctx, id, message); err != nil {
  37. return err
  38. }
  39. if account := r.accountsByID[id]; account != nil {
  40. account.Status = StatusError
  41. account.Schedulable = false
  42. account.ErrorMessage = message
  43. }
  44. return nil
  45. }
  46. type grokCredentialProxyRepoStub struct {
  47. ProxyRepository
  48. proxy *Proxy
  49. err error
  50. }
  51. func (r *grokCredentialProxyRepoStub) GetByID(context.Context, int64) (*Proxy, error) {
  52. return r.proxy, r.err
  53. }
  54. type grokCredentialBlockingRepo struct {
  55. *tokenRefreshAccountRepo
  56. setErrorStarted chan struct{}
  57. setTempStarted chan struct{}
  58. onceError sync.Once
  59. onceTemp sync.Once
  60. }
  61. type grokCredentialCommitThenCancelRepo struct {
  62. *tokenRefreshAccountRepo
  63. returnErr error
  64. }
  65. type grokCredentialUncommittedDeadlineRepo struct {
  66. *tokenRefreshAccountRepo
  67. }
  68. func (r *grokCredentialUncommittedDeadlineRepo) SetGrokCredentialErrorIfMatch(
  69. context.Context,
  70. int64,
  71. GrokCredentialMutationSnapshot,
  72. string,
  73. ) (bool, error) {
  74. return false, context.DeadlineExceeded
  75. }
  76. func (r *grokCredentialUncommittedDeadlineRepo) SetGrokCredentialTempUnschedulableIfMatch(
  77. context.Context,
  78. int64,
  79. GrokCredentialMutationSnapshot,
  80. time.Time,
  81. string,
  82. ) (bool, error) {
  83. return false, context.DeadlineExceeded
  84. }
  85. func (r *grokCredentialCommitThenCancelRepo) SetGrokCredentialErrorIfMatch(
  86. ctx context.Context,
  87. id int64,
  88. _ GrokCredentialMutationSnapshot,
  89. reason string,
  90. ) (bool, error) {
  91. account := r.accountsByID[id]
  92. account.Status = StatusError
  93. account.Schedulable = false
  94. account.ErrorMessage = reason
  95. if r.returnErr != nil {
  96. return false, r.returnErr
  97. }
  98. <-ctx.Done()
  99. return false, ctx.Err()
  100. }
  101. func (r *grokCredentialCommitThenCancelRepo) SetGrokCredentialTempUnschedulableIfMatch(
  102. ctx context.Context,
  103. id int64,
  104. _ GrokCredentialMutationSnapshot,
  105. until time.Time,
  106. reason string,
  107. ) (bool, error) {
  108. account := r.accountsByID[id]
  109. account.TempUnschedulableUntil = &until
  110. account.TempUnschedulableReason = reason
  111. if r.returnErr != nil {
  112. return false, r.returnErr
  113. }
  114. <-ctx.Done()
  115. return false, ctx.Err()
  116. }
  117. func (r *grokCredentialBlockingRepo) SetError(ctx context.Context, _ int64, _ string) error {
  118. r.onceError.Do(func() { close(r.setErrorStarted) })
  119. <-ctx.Done()
  120. return ctx.Err()
  121. }
  122. func (r *grokCredentialBlockingRepo) SetTempUnschedulable(ctx context.Context, _ int64, _ time.Time, _ string) error {
  123. r.onceTemp.Do(func() { close(r.setTempStarted) })
  124. <-ctx.Done()
  125. return ctx.Err()
  126. }
  127. func (r *grokCredentialBlockingRepo) SetGrokCredentialErrorIfMatch(
  128. ctx context.Context,
  129. _ int64,
  130. _ GrokCredentialMutationSnapshot,
  131. _ string,
  132. ) (bool, error) {
  133. r.onceError.Do(func() { close(r.setErrorStarted) })
  134. <-ctx.Done()
  135. return false, ctx.Err()
  136. }
  137. func (r *grokCredentialBlockingRepo) SetGrokCredentialTempUnschedulableIfMatch(
  138. ctx context.Context,
  139. _ int64,
  140. _ GrokCredentialMutationSnapshot,
  141. _ time.Time,
  142. _ string,
  143. ) (bool, error) {
  144. r.onceTemp.Do(func() { close(r.setTempStarted) })
  145. <-ctx.Done()
  146. return false, ctx.Err()
  147. }
  148. type grokCredentialBlockingCache struct {
  149. GrokTokenCache
  150. deleteStarted chan struct{}
  151. releaseDelete chan struct{}
  152. once sync.Once
  153. mu sync.Mutex
  154. deleted bool
  155. }
  156. type grokCredentialSequencedRepo struct {
  157. *tokenRefreshAccountRepo
  158. mu sync.Mutex
  159. latest *Account
  160. getCall int
  161. }
  162. type grokCredentialRereadFailureRepo struct {
  163. *tokenRefreshAccountRepo
  164. account *Account
  165. err error
  166. }
  167. type grokCredentialCountingRefresher struct {
  168. refreshCalls int
  169. }
  170. func (r *grokCredentialCountingRefresher) CacheKey(account *Account) string {
  171. return GrokTokenCacheKey(account)
  172. }
  173. func (r *grokCredentialCountingRefresher) CanRefresh(*Account) bool { return true }
  174. func (r *grokCredentialCountingRefresher) NeedsRefresh(*Account, time.Duration) bool { return true }
  175. func (r *grokCredentialCountingRefresher) Refresh(context.Context, *Account) (map[string]any, error) {
  176. r.refreshCalls++
  177. return map[string]any{"access_token": "must-not-be-used"}, nil
  178. }
  179. func (r *grokCredentialRereadFailureRepo) GetByID(context.Context, int64) (*Account, error) {
  180. return r.account, r.err
  181. }
  182. func (r *grokCredentialSequencedRepo) GetByID(ctx context.Context, id int64) (*Account, error) {
  183. r.mu.Lock()
  184. defer r.mu.Unlock()
  185. r.getCall++
  186. if r.getCall > 1 && r.latest != nil {
  187. return r.latest, nil
  188. }
  189. return r.tokenRefreshAccountRepo.GetByID(ctx, id)
  190. }
  191. func (c *grokCredentialBlockingCache) DeleteAccessToken(ctx context.Context, _ string) error {
  192. c.once.Do(func() { close(c.deleteStarted) })
  193. select {
  194. case <-c.releaseDelete:
  195. c.mu.Lock()
  196. c.deleted = true
  197. c.mu.Unlock()
  198. return nil
  199. case <-ctx.Done():
  200. return ctx.Err()
  201. }
  202. }
  203. func (c *grokCredentialBlockingCache) wasDeleted() bool {
  204. c.mu.Lock()
  205. defer c.mu.Unlock()
  206. return c.deleted
  207. }
  208. func TestUpstreamFailoverErrorNextAccountActionPreservesLegacyRetry(t *testing.T) {
  209. t.Parallel()
  210. require.True(t, (&UpstreamFailoverError{}).ShouldRetryNextAccount())
  211. require.True(t, (&UpstreamFailoverError{NextAccountAction: NextAccountRetry}).ShouldRetryNextAccount())
  212. require.False(t, (&UpstreamFailoverError{NextAccountAction: NextAccountStop}).ShouldRetryNextAccount())
  213. }
  214. func TestGetRequestCredentialMapsPermanentGrokOAuthFailureAndRedactsSecrets(t *testing.T) {
  215. gin.SetMode(gin.TestMode)
  216. account := expiredGrokOAuthAccountForCredentialTest(701)
  217. repo := &tokenRefreshAccountRepo{}
  218. repo.accountsByID = map[int64]*Account{account.ID: account}
  219. cache := &grokTokenCacheForProviderTest{lockResult: true}
  220. provider := NewGrokTokenProvider(repo, cache)
  221. provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{
  222. err: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_grant access_token=leaked-access refresh_token=leaked-refresh"),
  223. })
  224. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  225. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  226. token, kind, err := svc.getRequestCredential(context.Background(), c, account)
  227. require.Error(t, err)
  228. require.Empty(t, token)
  229. require.Empty(t, kind)
  230. var failoverErr *UpstreamFailoverError
  231. require.ErrorAs(t, err, &failoverErr)
  232. require.Equal(t, GatewayFailureStageAccountAuth, failoverErr.Stage)
  233. require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope)
  234. require.Equal(t, GrokCredentialReasonRevoked, failoverErr.Reason)
  235. require.True(t, failoverErr.ShouldRetryNextAccount())
  236. require.Equal(t, 0, failoverErr.StatusCode)
  237. require.Equal(t, http.StatusServiceUnavailable, failoverErr.ClientStatusCode)
  238. require.NotContains(t, err.Error(), "leaked-access")
  239. require.NotContains(t, err.Error(), "leaked-refresh")
  240. require.Equal(t, 1, repo.setErrorCalls)
  241. require.Zero(t, repo.setTempUnschedCalls)
  242. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  243. require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys)
  244. require.NotContains(t, repo.lastErrorMessage, "leaked-access")
  245. require.NotContains(t, repo.lastErrorMessage, "leaked-refresh")
  246. rawEvents, ok := c.Get(OpsUpstreamErrorsKey)
  247. require.True(t, ok)
  248. events, ok := rawEvents.([]*OpsUpstreamErrorEvent)
  249. require.True(t, ok)
  250. require.Len(t, events, 1)
  251. require.Equal(t, string(GatewayFailureStageAccountAuth), events[0].Stage)
  252. require.Equal(t, string(GatewayFailureScopeAccount), events[0].Scope)
  253. require.Equal(t, string(GrokCredentialReasonRevoked), events[0].Reason)
  254. require.Zero(t, events[0].UpstreamStatusCode)
  255. require.NotContains(t, events[0].Message, "leaked-access")
  256. require.NotContains(t, events[0].Message, "leaked-refresh")
  257. }
  258. func TestNewGrokCredentialFailoverDoesNotAttributeInferenceProxy(t *testing.T) {
  259. gin.SetMode(gin.TestMode)
  260. proxyID := int64(43)
  261. account := &Account{
  262. ID: 701,
  263. Platform: PlatformGrok,
  264. ProxyID: &proxyID,
  265. Proxy: &Proxy{ID: proxyID, Name: "bound-proxy"},
  266. }
  267. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  268. err := (&OpenAIGatewayService{}).newGrokCredentialFailover(c, account, grokCredentialFailureClass{
  269. scope: GatewayFailureScopeAccount,
  270. reason: GrokCredentialReasonAccountChanged,
  271. action: NextAccountRetry,
  272. message: "credential unavailable",
  273. })
  274. require.Error(t, err)
  275. rawEvents, ok := c.Get(OpsUpstreamErrorsKey)
  276. require.True(t, ok)
  277. events, ok := rawEvents.([]*OpsUpstreamErrorEvent)
  278. require.True(t, ok)
  279. require.Len(t, events, 1)
  280. require.Nil(t, events[0].ProxyID)
  281. require.Equal(t, opsProxyNameUnknown, events[0].ProxyName)
  282. }
  283. func TestGetRequestCredentialPermanentMappingsPersistAndInvalidate(t *testing.T) {
  284. gin.SetMode(gin.TestMode)
  285. tests := []struct {
  286. name string
  287. prepare func(*Account)
  288. refreshErr error
  289. wantReason GatewayFailureReason
  290. cachedToken string
  291. }{
  292. {
  293. name: "missing refresh credential",
  294. prepare: func(account *Account) {
  295. delete(account.Credentials, "refresh_token")
  296. },
  297. wantReason: GrokCredentialReasonMissing,
  298. },
  299. {
  300. name: "missing access credential",
  301. prepare: func(account *Account) {
  302. delete(account.Credentials, "access_token")
  303. account.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339)
  304. },
  305. wantReason: GrokCredentialReasonMissing,
  306. cachedToken: "stale-cached-access",
  307. },
  308. {
  309. name: "explicit entitlement action required",
  310. prepare: func(*Account) {},
  311. refreshErr: infraerrors.New(http.StatusForbidden, "GROK_OAUTH_ENTITLEMENT_DENIED", "access_denied"),
  312. wantReason: GrokCredentialReasonEntitlement,
  313. },
  314. }
  315. for index, tt := range tests {
  316. t.Run(tt.name, func(t *testing.T) {
  317. account := expiredGrokOAuthAccountForCredentialTest(int64(720 + index))
  318. account.Status = StatusActive
  319. account.Schedulable = true
  320. tt.prepare(account)
  321. baseRepo := &tokenRefreshAccountRepo{}
  322. baseRepo.accountsByID = map[int64]*Account{account.ID: account}
  323. repo := &grokCredentialPersistingRepo{tokenRefreshAccountRepo: baseRepo}
  324. cache := &grokTokenCacheForProviderTest{lockResult: true, token: tt.cachedToken}
  325. provider := NewGrokTokenProvider(repo, cache)
  326. provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{err: tt.refreshErr})
  327. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  328. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  329. _, _, err := svc.getRequestCredential(context.Background(), c, account)
  330. var failoverErr *UpstreamFailoverError
  331. require.ErrorAs(t, err, &failoverErr)
  332. require.Equal(t, tt.wantReason, failoverErr.Reason)
  333. require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope)
  334. require.Equal(t, 1, baseRepo.setErrorCalls)
  335. require.Equal(t, StatusError, account.Status)
  336. require.False(t, account.Schedulable)
  337. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  338. require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys)
  339. })
  340. }
  341. }
  342. func TestGetRequestCredentialMissingAccessNeverRefreshesAndPermanentlyFailsOver(t *testing.T) {
  343. gin.SetMode(gin.TestMode)
  344. tests := []struct {
  345. name string
  346. expiresAt *time.Time
  347. }{
  348. {name: "expiry missing"},
  349. {name: "expired", expiresAt: func() *time.Time { value := time.Now().Add(-time.Minute); return &value }()},
  350. {name: "near expiry", expiresAt: func() *time.Time { value := time.Now().Add(30 * time.Minute); return &value }()},
  351. }
  352. for index, tt := range tests {
  353. t.Run(tt.name, func(t *testing.T) {
  354. account := expiredGrokOAuthAccountForCredentialTest(int64(760 + index))
  355. account.Schedulable = true
  356. delete(account.Credentials, "access_token")
  357. if tt.expiresAt == nil {
  358. delete(account.Credentials, "expires_at")
  359. } else {
  360. account.Credentials["expires_at"] = tt.expiresAt.UTC().Format(time.RFC3339)
  361. }
  362. baseRepo := &tokenRefreshAccountRepo{}
  363. baseRepo.accountsByID = map[int64]*Account{account.ID: account}
  364. repo := &grokCredentialPersistingRepo{tokenRefreshAccountRepo: baseRepo}
  365. cache := &grokTokenCacheForProviderTest{lockResult: true, token: "stale-cache-must-not-win"}
  366. refresher := &grokCredentialCountingRefresher{}
  367. provider := NewGrokTokenProvider(repo, cache)
  368. provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), refresher)
  369. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  370. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  371. token, kind, err := svc.getRequestCredential(context.Background(), c, account)
  372. var failoverErr *UpstreamFailoverError
  373. require.ErrorAs(t, err, &failoverErr)
  374. require.Empty(t, token)
  375. require.Empty(t, kind)
  376. require.Equal(t, GrokCredentialReasonMissing, failoverErr.Reason)
  377. require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
  378. require.Zero(t, refresher.refreshCalls, "structurally missing access credentials must not reach the token endpoint")
  379. require.Equal(t, 1, baseRepo.setErrorCalls)
  380. require.Equal(t, StatusError, account.Status)
  381. require.False(t, account.Schedulable)
  382. require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys)
  383. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  384. })
  385. }
  386. }
  387. func TestGetRequestCredentialWarmCachedAccessWithMissingRefreshPermanentlyFailsOver(t *testing.T) {
  388. account := expiredGrokOAuthAccountForCredentialTest(764)
  389. account.Credentials["access_token"] = "valid-access"
  390. account.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339)
  391. delete(account.Credentials, "refresh_token")
  392. baseRepo := &tokenRefreshAccountRepo{}
  393. baseRepo.accountsByID = map[int64]*Account{account.ID: account}
  394. repo := &grokCredentialPersistingRepo{tokenRefreshAccountRepo: baseRepo}
  395. cache := &grokTokenCacheForProviderTest{lockResult: true, token: "valid-access"}
  396. refresher := &grokCredentialCountingRefresher{}
  397. provider := NewGrokTokenProvider(repo, cache)
  398. provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), refresher)
  399. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  400. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  401. _, _, err := svc.getRequestCredential(context.Background(), c, account)
  402. var failoverErr *UpstreamFailoverError
  403. require.ErrorAs(t, err, &failoverErr)
  404. require.Equal(t, GrokCredentialReasonMissing, failoverErr.Reason)
  405. require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
  406. require.Zero(t, refresher.refreshCalls)
  407. require.Equal(t, 1, baseRepo.setErrorCalls)
  408. require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys)
  409. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  410. }
  411. func TestGetRequestCredentialMapsTransientAndProviderFailuresSeparately(t *testing.T) {
  412. gin.SetMode(gin.TestMode)
  413. t.Run("account transient temporarily unschedules", func(t *testing.T) {
  414. account := expiredGrokOAuthAccountForCredentialTest(702)
  415. repo := &tokenRefreshAccountRepo{}
  416. repo.accountsByID = map[int64]*Account{account.ID: account}
  417. cache := &grokTokenCacheForProviderTest{lockResult: true}
  418. provider := NewGrokTokenProvider(repo, cache)
  419. provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{err: errors.New("temporary refresh transport failure")})
  420. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  421. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  422. _, _, err := svc.getRequestCredential(context.Background(), c, account)
  423. var failoverErr *UpstreamFailoverError
  424. require.ErrorAs(t, err, &failoverErr)
  425. require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope)
  426. require.Equal(t, GrokCredentialReasonRefreshTransient, failoverErr.Reason)
  427. require.True(t, failoverErr.ShouldRetryNextAccount())
  428. require.Zero(t, repo.setErrorCalls)
  429. require.Equal(t, 1, repo.setTempUnschedCalls)
  430. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  431. })
  432. t.Run("shared provider configuration stops without mutation", func(t *testing.T) {
  433. account := expiredGrokOAuthAccountForCredentialTest(703)
  434. repo := &tokenRefreshAccountRepo{}
  435. repo.accountsByID = map[int64]*Account{account.ID: account}
  436. provider := NewGrokTokenProvider(repo, nil)
  437. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  438. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  439. _, _, err := svc.getRequestCredential(context.Background(), c, account)
  440. var failoverErr *UpstreamFailoverError
  441. require.ErrorAs(t, err, &failoverErr)
  442. require.Equal(t, GatewayFailureScopeProvider, failoverErr.Scope)
  443. require.Equal(t, NextAccountStop, failoverErr.NextAccountAction)
  444. require.Zero(t, repo.setErrorCalls)
  445. require.Zero(t, repo.setTempUnschedCalls)
  446. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  447. })
  448. t.Run("account reread failures preserve shared versus missing-row scope", func(t *testing.T) {
  449. for _, tt := range []struct {
  450. name string
  451. account *Account
  452. err error
  453. }{
  454. {name: "repository error", err: errors.New("database temporarily unavailable")},
  455. {name: "missing row"},
  456. } {
  457. t.Run(tt.name, func(t *testing.T) {
  458. account := expiredGrokOAuthAccountForCredentialTest(712)
  459. baseRepo := &tokenRefreshAccountRepo{}
  460. baseRepo.accountsByID = map[int64]*Account{account.ID: account}
  461. repo := &grokCredentialRereadFailureRepo{tokenRefreshAccountRepo: baseRepo, account: tt.account, err: tt.err}
  462. cache := &grokTokenCacheForProviderTest{lockResult: true}
  463. provider := NewGrokTokenProvider(repo, cache)
  464. provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), NewGrokTokenRefresher(nil))
  465. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  466. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  467. _, _, err := svc.getRequestCredential(context.Background(), c, account)
  468. var failoverErr *UpstreamFailoverError
  469. require.ErrorAs(t, err, &failoverErr)
  470. if tt.err != nil {
  471. require.Equal(t, GatewayFailureScopeProvider, failoverErr.Scope)
  472. require.Equal(t, GrokCredentialReasonProviderDown, failoverErr.Reason)
  473. require.Equal(t, NextAccountStop, failoverErr.NextAccountAction)
  474. } else {
  475. require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope)
  476. require.Equal(t, GrokCredentialReasonAccountChanged, failoverErr.Reason)
  477. require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
  478. }
  479. require.Zero(t, baseRepo.setErrorCalls)
  480. require.Zero(t, baseRepo.setTempUnschedCalls)
  481. require.Empty(t, cache.deletedKeys)
  482. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  483. })
  484. }
  485. })
  486. t.Run("fresh account eligibility changes retry without mutating stale state", func(t *testing.T) {
  487. for _, tt := range []struct {
  488. name string
  489. mutate func(*Account)
  490. }{
  491. {
  492. name: "account disabled",
  493. mutate: func(account *Account) {
  494. account.Status = StatusDisabled
  495. },
  496. },
  497. {
  498. name: "account converted",
  499. mutate: func(account *Account) {
  500. account.Type = AccountTypeUpstream
  501. },
  502. },
  503. {
  504. name: "account manually unschedulable",
  505. mutate: func(account *Account) {
  506. account.Schedulable = false
  507. },
  508. },
  509. {
  510. name: "account temporarily unschedulable",
  511. mutate: func(account *Account) {
  512. until := time.Now().Add(time.Minute)
  513. account.TempUnschedulableUntil = &until
  514. },
  515. },
  516. } {
  517. t.Run(tt.name, func(t *testing.T) {
  518. staleAccount := expiredGrokOAuthAccountForCredentialTest(713)
  519. freshAccount := *staleAccount
  520. freshAccount.Credentials = shallowCopyMap(staleAccount.Credentials)
  521. tt.mutate(&freshAccount)
  522. repo := &tokenRefreshAccountRepo{}
  523. repo.accountsByID = map[int64]*Account{staleAccount.ID: &freshAccount}
  524. cache := &grokTokenCacheForProviderTest{lockResult: true}
  525. provider := NewGrokTokenProvider(repo, cache)
  526. provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), NewGrokTokenRefresher(nil))
  527. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  528. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  529. _, _, err := svc.getRequestCredential(context.Background(), c, staleAccount)
  530. var failoverErr *UpstreamFailoverError
  531. require.ErrorAs(t, err, &failoverErr)
  532. require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope)
  533. require.Equal(t, GrokCredentialReasonAccountChanged, failoverErr.Reason)
  534. require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
  535. require.Zero(t, repo.setErrorCalls)
  536. require.Zero(t, repo.setTempUnschedCalls)
  537. require.Empty(t, cache.deletedKeys)
  538. require.False(t, svc.isOpenAIAccountRuntimeBlocked(staleAccount))
  539. })
  540. }
  541. })
  542. t.Run("fresh missing refresh credential permanently blocks the account", func(t *testing.T) {
  543. staleAccount := expiredGrokOAuthAccountForCredentialTest(714)
  544. staleAccount.Schedulable = true
  545. freshAccount := *staleAccount
  546. freshAccount.Credentials = shallowCopyMap(staleAccount.Credentials)
  547. delete(freshAccount.Credentials, "refresh_token")
  548. baseRepo := &tokenRefreshAccountRepo{}
  549. baseRepo.accountsByID = map[int64]*Account{staleAccount.ID: &freshAccount}
  550. repo := &grokCredentialPersistingRepo{tokenRefreshAccountRepo: baseRepo}
  551. cache := &grokTokenCacheForProviderTest{lockResult: true}
  552. provider := NewGrokTokenProvider(repo, cache)
  553. provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), NewGrokTokenRefresher(nil))
  554. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  555. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  556. _, _, err := svc.getRequestCredential(context.Background(), c, staleAccount)
  557. var failoverErr *UpstreamFailoverError
  558. require.ErrorAs(t, err, &failoverErr)
  559. require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope)
  560. require.Equal(t, GrokCredentialReasonMissing, failoverErr.Reason)
  561. require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
  562. require.Equal(t, 1, baseRepo.setErrorCalls)
  563. require.Zero(t, baseRepo.setTempUnschedCalls)
  564. require.Equal(t, StatusError, freshAccount.Status)
  565. require.False(t, freshAccount.Schedulable)
  566. require.Equal(t, []string{GrokTokenCacheKey(staleAccount)}, cache.deletedKeys)
  567. require.True(t, svc.isOpenAIAccountRuntimeBlocked(staleAccount))
  568. })
  569. t.Run("refresh added after locked structural failure wins conditional mutation", func(t *testing.T) {
  570. staleAccount := expiredGrokOAuthAccountForCredentialTest(717)
  571. freshAccount := *staleAccount
  572. freshAccount.Credentials = shallowCopyMap(staleAccount.Credentials)
  573. delete(freshAccount.Credentials, "refresh_token")
  574. repo := &tokenRefreshAccountRepo{}
  575. repo.accountsByID = map[int64]*Account{staleAccount.ID: &freshAccount}
  576. repo.beforeConditionalState = func() {
  577. repaired := freshAccount
  578. repaired.Credentials = shallowCopyMap(freshAccount.Credentials)
  579. repaired.Credentials["refresh_token"] = "repaired-refresh-token"
  580. repaired.Credentials["expires_at"] = time.Now().Add(time.Hour).UTC().Format(time.RFC3339)
  581. repo.accountsByID[staleAccount.ID] = &repaired
  582. }
  583. cache := &grokTokenCacheForProviderTest{lockResult: true}
  584. provider := NewGrokTokenProvider(repo, cache)
  585. provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), NewGrokTokenRefresher(nil))
  586. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  587. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  588. token, kind, err := svc.getRequestCredential(context.Background(), c, staleAccount)
  589. require.NoError(t, err)
  590. require.Equal(t, "expired-access-token", token)
  591. require.Equal(t, "oauth", kind)
  592. require.Zero(t, repo.setErrorCalls)
  593. require.Zero(t, repo.setTempUnschedCalls)
  594. require.Empty(t, cache.deletedKeys)
  595. require.False(t, svc.isOpenAIAccountRuntimeBlocked(staleAccount))
  596. })
  597. t.Run("expiry-only repair wins full credential fingerprint CAS", func(t *testing.T) {
  598. account := expiredGrokOAuthAccountForCredentialTest(718)
  599. repo := &tokenRefreshAccountRepo{}
  600. repo.accountsByID = map[int64]*Account{account.ID: account}
  601. repo.beforeConditionalState = func() {
  602. repaired := *account
  603. repaired.Credentials = shallowCopyMap(account.Credentials)
  604. repaired.Credentials["expires_at"] = time.Now().Add(time.Hour).UTC().Format(time.RFC3339)
  605. repo.accountsByID[account.ID] = &repaired
  606. }
  607. cache := &grokTokenCacheForProviderTest{lockResult: true}
  608. provider := NewGrokTokenProvider(repo, cache)
  609. provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{
  610. err: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_grant"),
  611. })
  612. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  613. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  614. token, kind, err := svc.getRequestCredential(context.Background(), c, account)
  615. require.NoError(t, err)
  616. require.Equal(t, "expired-access-token", token)
  617. require.Equal(t, "oauth", kind)
  618. require.Zero(t, repo.setErrorCalls)
  619. require.Empty(t, cache.deletedKeys)
  620. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  621. })
  622. t.Run("generic token endpoint 403 stops as shared provider failure", func(t *testing.T) {
  623. account := expiredGrokOAuthAccountForCredentialTest(708)
  624. repo := &tokenRefreshAccountRepo{}
  625. repo.accountsByID = map[int64]*Account{account.ID: account}
  626. cache := &grokTokenCacheForProviderTest{lockResult: true}
  627. provider := NewGrokTokenProvider(repo, cache)
  628. provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{
  629. err: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "token refresh failed: status 403, body: forbidden"),
  630. })
  631. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  632. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  633. _, _, err := svc.getRequestCredential(context.Background(), c, account)
  634. var failoverErr *UpstreamFailoverError
  635. require.ErrorAs(t, err, &failoverErr)
  636. require.Equal(t, GatewayFailureScopeProvider, failoverErr.Scope)
  637. require.Equal(t, GrokCredentialReasonProviderDown, failoverErr.Reason)
  638. require.Equal(t, NextAccountStop, failoverErr.NextAccountAction)
  639. require.Zero(t, repo.setErrorCalls)
  640. require.Zero(t, repo.setTempUnschedCalls)
  641. require.Empty(t, cache.deletedKeys)
  642. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  643. })
  644. t.Run("account proxy generic 403 remains bounded account transient", func(t *testing.T) {
  645. account := expiredGrokOAuthAccountForCredentialTest(711)
  646. proxyID := int64(43)
  647. account.ProxyID = &proxyID
  648. account.Proxy = &Proxy{}
  649. repo := &tokenRefreshAccountRepo{}
  650. repo.accountsByID = map[int64]*Account{account.ID: account}
  651. cache := &grokTokenCacheForProviderTest{lockResult: true}
  652. provider := NewGrokTokenProvider(repo, cache)
  653. provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{
  654. err: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "token refresh failed: status 403, body: forbidden"),
  655. })
  656. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  657. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  658. _, _, err := svc.getRequestCredential(context.Background(), c, account)
  659. var failoverErr *UpstreamFailoverError
  660. require.ErrorAs(t, err, &failoverErr)
  661. require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope)
  662. require.Equal(t, GrokCredentialReasonRefreshTransient, failoverErr.Reason)
  663. require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
  664. require.Zero(t, repo.setErrorCalls)
  665. require.Equal(t, 1, repo.setTempUnschedCalls)
  666. require.Empty(t, cache.deletedKeys)
  667. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  668. })
  669. t.Run("proxy repository read failure stops without account mutation", func(t *testing.T) {
  670. account := expiredGrokOAuthAccountForCredentialTest(709)
  671. proxyID := int64(41)
  672. account.ProxyID = &proxyID
  673. account.Proxy = &Proxy{}
  674. repo := &tokenRefreshAccountRepo{}
  675. repo.accountsByID = map[int64]*Account{account.ID: account}
  676. cache := &grokTokenCacheForProviderTest{lockResult: true}
  677. oauthSvc := NewGrokOAuthService(&grokCredentialProxyRepoStub{err: errors.New("database temporarily unavailable")}, &grokOAuthClientStub{})
  678. defer oauthSvc.Stop()
  679. provider := NewGrokTokenProvider(repo, cache)
  680. provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), NewGrokTokenRefresher(oauthSvc))
  681. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  682. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  683. _, _, err := svc.getRequestCredential(context.Background(), c, account)
  684. var failoverErr *UpstreamFailoverError
  685. require.ErrorAs(t, err, &failoverErr)
  686. require.Equal(t, GatewayFailureScopeProvider, failoverErr.Scope)
  687. require.Equal(t, GrokCredentialReasonProviderDown, failoverErr.Reason)
  688. require.Equal(t, NextAccountStop, failoverErr.NextAccountAction)
  689. require.Zero(t, repo.setErrorCalls)
  690. require.Zero(t, repo.setTempUnschedCalls)
  691. require.Empty(t, cache.deletedKeys)
  692. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  693. })
  694. t.Run("structurally missing configured proxy permanently blocks only that account", func(t *testing.T) {
  695. account := expiredGrokOAuthAccountForCredentialTest(710)
  696. account.Status = StatusActive
  697. account.Schedulable = true
  698. proxyID := int64(42)
  699. account.ProxyID = &proxyID
  700. baseRepo := &tokenRefreshAccountRepo{}
  701. baseRepo.accountsByID = map[int64]*Account{account.ID: account}
  702. repo := &grokCredentialPersistingRepo{tokenRefreshAccountRepo: baseRepo}
  703. cache := &grokTokenCacheForProviderTest{lockResult: true}
  704. oauthSvc := NewGrokOAuthService(&grokCredentialProxyRepoStub{err: ErrProxyNotFound}, &grokOAuthClientStub{})
  705. defer oauthSvc.Stop()
  706. provider := NewGrokTokenProvider(repo, cache)
  707. provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), NewGrokTokenRefresher(oauthSvc))
  708. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  709. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  710. _, _, err := svc.getRequestCredential(context.Background(), c, account)
  711. var failoverErr *UpstreamFailoverError
  712. require.ErrorAs(t, err, &failoverErr)
  713. require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope)
  714. require.Equal(t, GrokCredentialReasonProxyInvalid, failoverErr.Reason)
  715. require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
  716. require.Equal(t, 1, baseRepo.setErrorCalls)
  717. require.Equal(t, StatusError, account.Status)
  718. require.False(t, account.Schedulable)
  719. require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys)
  720. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  721. })
  722. }
  723. func TestGetRequestCredentialRuntimeBlockWinsBeforeWarmTokenCache(t *testing.T) {
  724. account := expiredGrokOAuthAccountForCredentialTest(716)
  725. account.Credentials["access_token"] = "valid-access"
  726. account.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339)
  727. repo := &tokenRefreshAccountRepo{}
  728. repo.accountsByID = map[int64]*Account{account.ID: account}
  729. cache := &grokTokenCacheForProviderTest{token: "valid-access"}
  730. provider := NewGrokTokenProvider(repo, cache)
  731. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  732. svc.BlockAccountScheduling(account, time.Now().Add(time.Minute), "independent")
  733. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  734. _, _, err := svc.getRequestCredential(context.Background(), c, account)
  735. var failoverErr *UpstreamFailoverError
  736. require.ErrorAs(t, err, &failoverErr)
  737. require.Equal(t, GrokCredentialReasonAccountChanged, failoverErr.Reason)
  738. require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
  739. require.Zero(t, cache.getCalls)
  740. require.Zero(t, repo.setErrorCalls)
  741. require.Zero(t, repo.setTempUnschedCalls)
  742. }
  743. func TestGetRequestCredentialWarmCachedAccessWithMissingConfiguredProxyPermanentlyFailsOver(t *testing.T) {
  744. account := expiredGrokOAuthAccountForCredentialTest(715)
  745. account.Credentials["access_token"] = "valid-access"
  746. account.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339)
  747. proxyID := int64(44)
  748. account.ProxyID = &proxyID
  749. account.Proxy = nil
  750. baseRepo := &tokenRefreshAccountRepo{}
  751. baseRepo.accountsByID = map[int64]*Account{account.ID: account}
  752. repo := &grokCredentialPersistingRepo{tokenRefreshAccountRepo: baseRepo}
  753. cache := &grokTokenCacheForProviderTest{lockResult: true, token: "valid-access"}
  754. refresher := &grokCredentialCountingRefresher{}
  755. provider := NewGrokTokenProvider(repo, cache)
  756. provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), refresher)
  757. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  758. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  759. _, _, err := svc.getRequestCredential(context.Background(), c, account)
  760. var failoverErr *UpstreamFailoverError
  761. require.ErrorAs(t, err, &failoverErr)
  762. require.Equal(t, GrokCredentialReasonProxyInvalid, failoverErr.Reason)
  763. require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
  764. require.Zero(t, refresher.refreshCalls)
  765. require.Equal(t, 1, baseRepo.setErrorCalls)
  766. require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys)
  767. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  768. }
  769. func TestGetRequestCredentialCancellationAndBudgetDoNotMutateAccount(t *testing.T) {
  770. gin.SetMode(gin.TestMode)
  771. account := expiredGrokOAuthAccountForCredentialTest(704)
  772. repo := &tokenRefreshAccountRepo{}
  773. repo.accountsByID = map[int64]*Account{account.ID: account}
  774. provider := NewGrokTokenProvider(repo, nil)
  775. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  776. t.Run("parent cancellation is returned directly", func(t *testing.T) {
  777. ctx, cancel := context.WithCancel(context.Background())
  778. cancel()
  779. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  780. _, _, err := svc.getRequestCredential(ctx, c, account)
  781. require.ErrorIs(t, err, context.Canceled)
  782. var failoverErr *UpstreamFailoverError
  783. require.False(t, errors.As(err, &failoverErr))
  784. })
  785. t.Run("request credential budget stops safely", func(t *testing.T) {
  786. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  787. c.Set(grokCredentialFailoverDeadlineKey, time.Now().Add(-time.Second))
  788. _, _, err := svc.getRequestCredential(context.Background(), c, account)
  789. var failoverErr *UpstreamFailoverError
  790. require.ErrorAs(t, err, &failoverErr)
  791. require.Equal(t, GatewayFailureScopeRequest, failoverErr.Scope)
  792. require.Equal(t, GrokCredentialReasonFailoverTimeout, failoverErr.Reason)
  793. require.False(t, failoverErr.ShouldRetryNextAccount())
  794. })
  795. require.Zero(t, repo.setErrorCalls)
  796. require.Zero(t, repo.setTempUnschedCalls)
  797. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  798. }
  799. func TestGetRequestCredentialStateMutationFailureStopsAndKeepsRuntimeBlock(t *testing.T) {
  800. gin.SetMode(gin.TestMode)
  801. tests := []struct {
  802. name string
  803. refreshErr error
  804. configure func(*tokenRefreshAccountRepo, *grokTokenCacheForProviderTest)
  805. wantSetError int
  806. wantSetTemp int
  807. wantCacheDelete int
  808. }{
  809. {
  810. name: "permanent state persistence",
  811. refreshErr: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_grant"),
  812. configure: func(repo *tokenRefreshAccountRepo, _ *grokTokenCacheForProviderTest) {
  813. repo.setErrorErr = errors.New("database write failed")
  814. },
  815. wantSetError: 1,
  816. },
  817. {
  818. name: "transient state persistence",
  819. refreshErr: errors.New("temporary refresh transport failure"),
  820. configure: func(repo *tokenRefreshAccountRepo, _ *grokTokenCacheForProviderTest) {
  821. repo.setTempUnschedErr = errors.New("database write failed")
  822. },
  823. wantSetTemp: 1,
  824. },
  825. {
  826. name: "permanent token cache invalidation",
  827. refreshErr: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_grant"),
  828. configure: func(_ *tokenRefreshAccountRepo, cache *grokTokenCacheForProviderTest) {
  829. cache.deleteErr = errors.New("cache delete failed")
  830. },
  831. wantSetError: 1,
  832. wantCacheDelete: 1,
  833. },
  834. }
  835. for index, tt := range tests {
  836. t.Run(tt.name, func(t *testing.T) {
  837. account := expiredGrokOAuthAccountForCredentialTest(int64(740 + index))
  838. repo := &tokenRefreshAccountRepo{}
  839. repo.accountsByID = map[int64]*Account{account.ID: account}
  840. cache := &grokTokenCacheForProviderTest{lockResult: true}
  841. tt.configure(repo, cache)
  842. provider := NewGrokTokenProvider(repo, cache)
  843. provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{err: tt.refreshErr})
  844. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  845. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  846. _, _, err := svc.getRequestCredential(context.Background(), c, account)
  847. var failoverErr *UpstreamFailoverError
  848. require.ErrorAs(t, err, &failoverErr)
  849. require.Equal(t, GatewayFailureScopeProvider, failoverErr.Scope)
  850. require.Equal(t, GrokCredentialReasonStateUpdate, failoverErr.Reason)
  851. require.Equal(t, NextAccountStop, failoverErr.NextAccountAction)
  852. require.Equal(t, tt.wantSetError, repo.setErrorCalls)
  853. require.Equal(t, tt.wantSetTemp, repo.setTempUnschedCalls)
  854. require.Len(t, cache.deletedKeys, tt.wantCacheDelete)
  855. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account), "failed mutation must retain the immediate local block")
  856. })
  857. }
  858. }
  859. func TestGrokCredentialMutationBoundariesHonorParentCancellation(t *testing.T) {
  860. t.Run("blocked SetError cancellation prevents cache and runtime mutation", func(t *testing.T) {
  861. account := expiredGrokOAuthAccountForCredentialTest(730)
  862. baseRepo := &tokenRefreshAccountRepo{}
  863. baseRepo.accountsByID = map[int64]*Account{account.ID: account}
  864. repo := &grokCredentialBlockingRepo{
  865. tokenRefreshAccountRepo: baseRepo,
  866. setErrorStarted: make(chan struct{}),
  867. setTempStarted: make(chan struct{}),
  868. }
  869. cache := &grokTokenCacheForProviderTest{}
  870. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, cache)}
  871. ctx, cancel := context.WithCancel(context.Background())
  872. result := make(chan error, 1)
  873. go func() {
  874. _, err := svc.applyGrokCredentialAccountFailure(ctx, account, grokCredentialFailureClass{
  875. reason: GrokCredentialReasonRevoked, permanent: true,
  876. })
  877. result <- err
  878. }()
  879. <-repo.setErrorStarted
  880. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account), "runtime block must precede persistent SetError")
  881. cancel()
  882. require.ErrorIs(t, <-result, context.Canceled)
  883. require.Zero(t, baseRepo.setErrorCalls)
  884. require.Empty(t, cache.deletedKeys)
  885. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  886. })
  887. t.Run("blocked temporary unschedule cancellation prevents runtime mutation", func(t *testing.T) {
  888. account := expiredGrokOAuthAccountForCredentialTest(731)
  889. baseRepo := &tokenRefreshAccountRepo{}
  890. baseRepo.accountsByID = map[int64]*Account{account.ID: account}
  891. repo := &grokCredentialBlockingRepo{
  892. tokenRefreshAccountRepo: baseRepo,
  893. setErrorStarted: make(chan struct{}),
  894. setTempStarted: make(chan struct{}),
  895. }
  896. svc := &OpenAIGatewayService{accountRepo: repo}
  897. ctx, cancel := context.WithCancel(context.Background())
  898. result := make(chan error, 1)
  899. go func() {
  900. _, err := svc.applyGrokCredentialAccountFailure(ctx, account, grokCredentialFailureClass{
  901. reason: GrokCredentialReasonRefreshTransient, transient: true,
  902. })
  903. result <- err
  904. }()
  905. <-repo.setTempStarted
  906. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account), "runtime block must precede temporary unscheduling")
  907. cancel()
  908. require.ErrorIs(t, <-result, context.Canceled)
  909. require.Zero(t, baseRepo.setTempUnschedCalls)
  910. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  911. })
  912. t.Run("post-commit cancellation finishes cache cleanup and retains quarantine", func(t *testing.T) {
  913. account := expiredGrokOAuthAccountForCredentialTest(732)
  914. repo := &tokenRefreshAccountRepo{}
  915. repo.accountsByID = map[int64]*Account{account.ID: account}
  916. cache := &grokCredentialBlockingCache{deleteStarted: make(chan struct{}), releaseDelete: make(chan struct{})}
  917. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, cache)}
  918. ctx, cancel := context.WithCancel(context.Background())
  919. result := make(chan error, 1)
  920. go func() {
  921. _, err := svc.applyGrokCredentialAccountFailure(ctx, account, grokCredentialFailureClass{
  922. reason: GrokCredentialReasonRevoked, permanent: true,
  923. })
  924. result <- err
  925. }()
  926. <-cache.deleteStarted
  927. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account), "runtime block must precede cache invalidation")
  928. cancel()
  929. close(cache.releaseDelete)
  930. require.ErrorIs(t, <-result, context.Canceled)
  931. require.Equal(t, 1, repo.setErrorCalls)
  932. require.True(t, cache.wasDeleted())
  933. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  934. })
  935. }
  936. func TestGrokCredentialMutationLockWaitHonorsCredentialBudget(t *testing.T) {
  937. account := expiredGrokOAuthAccountForCredentialTest(735)
  938. repo := &tokenRefreshAccountRepo{}
  939. repo.accountsByID = map[int64]*Account{account.ID: account}
  940. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, &grokTokenCacheForProviderTest{})}
  941. mutationLock := svc.grokCredentialMutationLock(account.ID)
  942. require.NoError(t, mutationLock.Lock(context.Background()))
  943. defer mutationLock.Unlock()
  944. ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
  945. defer cancel()
  946. startedAt := time.Now()
  947. token, err := svc.applyGrokCredentialAccountFailure(ctx, account, grokCredentialFailureClass{
  948. reason: GrokCredentialReasonRevoked, permanent: true,
  949. })
  950. require.ErrorIs(t, err, context.DeadlineExceeded)
  951. require.Empty(t, token)
  952. require.Less(t, time.Since(startedAt), 500*time.Millisecond)
  953. require.Zero(t, repo.setErrorCalls)
  954. require.Zero(t, repo.setTempUnschedCalls)
  955. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  956. }
  957. func TestGetRequestCredentialBudgetBoundsBlockedConditionalMutation(t *testing.T) {
  958. gin.SetMode(gin.TestMode)
  959. account := expiredGrokOAuthAccountForCredentialTest(736)
  960. baseRepo := &tokenRefreshAccountRepo{}
  961. baseRepo.accountsByID = map[int64]*Account{account.ID: account}
  962. repo := &grokCredentialBlockingRepo{
  963. tokenRefreshAccountRepo: baseRepo,
  964. setErrorStarted: make(chan struct{}),
  965. setTempStarted: make(chan struct{}),
  966. }
  967. cache := &grokTokenCacheForProviderTest{lockResult: true}
  968. provider := NewGrokTokenProvider(repo, cache)
  969. provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{
  970. err: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_grant"),
  971. })
  972. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  973. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  974. c.Set(grokCredentialFailoverDeadlineKey, time.Now().Add(40*time.Millisecond))
  975. startedAt := time.Now()
  976. token, kind, err := svc.getRequestCredential(context.Background(), c, account)
  977. var failoverErr *UpstreamFailoverError
  978. require.ErrorAs(t, err, &failoverErr)
  979. require.Equal(t, GatewayFailureScopeRequest, failoverErr.Scope)
  980. require.Equal(t, GrokCredentialReasonFailoverTimeout, failoverErr.Reason)
  981. require.Equal(t, NextAccountStop, failoverErr.NextAccountAction)
  982. require.Empty(t, token)
  983. require.Empty(t, kind)
  984. require.Less(t, time.Since(startedAt), 500*time.Millisecond)
  985. require.Zero(t, baseRepo.setErrorCalls)
  986. require.Zero(t, baseRepo.setTempUnschedCalls)
  987. require.Empty(t, cache.deletedKeys)
  988. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  989. }
  990. func TestGetRequestCredentialLockHeldTimeoutDoesNotQuarantineAccount(t *testing.T) {
  991. gin.SetMode(gin.TestMode)
  992. tests := []struct {
  993. name string
  994. buildRepo func(*Account) AccountRepository
  995. wantScope GatewayFailureScope
  996. wantReason GatewayFailureReason
  997. wantAction NextAccountAction
  998. }{
  999. {
  1000. name: "authoritative row unchanged",
  1001. buildRepo: func(account *Account) AccountRepository {
  1002. repo := &tokenRefreshAccountRepo{}
  1003. repo.accountsByID = map[int64]*Account{account.ID: account}
  1004. return repo
  1005. },
  1006. wantScope: GatewayFailureScopeAccount,
  1007. wantReason: GrokCredentialReasonAccountChanged,
  1008. wantAction: NextAccountRetry,
  1009. },
  1010. {
  1011. name: "selected account was deleted",
  1012. buildRepo: func(account *Account) AccountRepository {
  1013. base := &tokenRefreshAccountRepo{}
  1014. base.accountsByID = map[int64]*Account{account.ID: account}
  1015. return &grokCredentialRereadFailureRepo{tokenRefreshAccountRepo: base}
  1016. },
  1017. wantScope: GatewayFailureScopeAccount,
  1018. wantReason: GrokCredentialReasonAccountChanged,
  1019. wantAction: NextAccountRetry,
  1020. },
  1021. {
  1022. name: "shared account store unavailable",
  1023. buildRepo: func(account *Account) AccountRepository {
  1024. base := &tokenRefreshAccountRepo{}
  1025. base.accountsByID = map[int64]*Account{account.ID: account}
  1026. return &grokCredentialRereadFailureRepo{tokenRefreshAccountRepo: base, err: errors.New("database unavailable")}
  1027. },
  1028. wantScope: GatewayFailureScopeProvider,
  1029. wantReason: GrokCredentialReasonProviderDown,
  1030. wantAction: NextAccountStop,
  1031. },
  1032. }
  1033. for index, tt := range tests {
  1034. t.Run(tt.name, func(t *testing.T) {
  1035. account := expiredGrokOAuthAccountForCredentialTest(int64(7400 + index))
  1036. repo := tt.buildRepo(account)
  1037. cache := &grokTokenCacheForProviderTest{lockResult: false}
  1038. provider := NewGrokTokenProvider(repo, cache)
  1039. provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{})
  1040. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  1041. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  1042. startedAt := time.Now()
  1043. _, _, err := svc.getRequestCredential(context.Background(), c, account)
  1044. var failoverErr *UpstreamFailoverError
  1045. require.ErrorAs(t, err, &failoverErr)
  1046. require.Equal(t, tt.wantScope, failoverErr.Scope)
  1047. require.Equal(t, tt.wantReason, failoverErr.Reason)
  1048. require.Equal(t, tt.wantAction, failoverErr.NextAccountAction)
  1049. require.Less(t, time.Since(startedAt), 3*time.Second)
  1050. switch countingRepo := repo.(type) {
  1051. case *tokenRefreshAccountRepo:
  1052. require.Zero(t, countingRepo.setErrorCalls)
  1053. require.Zero(t, countingRepo.setTempUnschedCalls)
  1054. case *grokCredentialRereadFailureRepo:
  1055. require.Zero(t, countingRepo.tokenRefreshAccountRepo.setErrorCalls)
  1056. require.Zero(t, countingRepo.tokenRefreshAccountRepo.setTempUnschedCalls)
  1057. }
  1058. require.Empty(t, cache.deletedKeys)
  1059. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  1060. })
  1061. }
  1062. }
  1063. func TestGrokCredentialMutationCancellationAmbiguityConfirmsDurableCommit(t *testing.T) {
  1064. tests := []struct {
  1065. name string
  1066. class grokCredentialFailureClass
  1067. committed func(*Account) bool
  1068. }{
  1069. {
  1070. name: "permanent quarantine",
  1071. class: grokCredentialFailureClass{reason: GrokCredentialReasonRevoked, permanent: true},
  1072. committed: func(account *Account) bool {
  1073. return account.Status == StatusError && !account.Schedulable && account.ErrorMessage == string(GrokCredentialReasonRevoked)
  1074. },
  1075. },
  1076. {
  1077. name: "temporary quarantine",
  1078. class: grokCredentialFailureClass{reason: GrokCredentialReasonRefreshTransient, transient: true},
  1079. committed: func(account *Account) bool {
  1080. return account.TempUnschedulableUntil != nil && account.TempUnschedulableReason == string(GrokCredentialReasonRefreshTransient)
  1081. },
  1082. },
  1083. }
  1084. for index, tt := range tests {
  1085. t.Run(tt.name, func(t *testing.T) {
  1086. account := expiredGrokOAuthAccountForCredentialTest(int64(737 + index))
  1087. baseRepo := &tokenRefreshAccountRepo{}
  1088. baseRepo.accountsByID = map[int64]*Account{account.ID: account}
  1089. repo := &grokCredentialCommitThenCancelRepo{tokenRefreshAccountRepo: baseRepo}
  1090. cache := &grokTokenCacheForProviderTest{}
  1091. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, cache)}
  1092. ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
  1093. defer cancel()
  1094. token, err := svc.applyGrokCredentialAccountFailure(ctx, account, tt.class)
  1095. require.ErrorIs(t, err, context.DeadlineExceeded)
  1096. require.Empty(t, token)
  1097. require.True(t, tt.committed(account), "the detached confirmation must recognize the durable mutation")
  1098. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account), "a confirmed durable quarantine must retain its runtime block")
  1099. if tt.class.permanent {
  1100. require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys)
  1101. } else {
  1102. require.Empty(t, cache.deletedKeys)
  1103. }
  1104. })
  1105. }
  1106. }
  1107. func TestGrokCredentialInnerStateDeadlineAmbiguityConfirmsDurableCommit(t *testing.T) {
  1108. tests := []struct {
  1109. name string
  1110. class grokCredentialFailureClass
  1111. }{
  1112. {name: "permanent quarantine", class: grokCredentialFailureClass{reason: GrokCredentialReasonRevoked, permanent: true}},
  1113. {name: "temporary quarantine", class: grokCredentialFailureClass{reason: GrokCredentialReasonRefreshTransient, transient: true}},
  1114. }
  1115. for index, tt := range tests {
  1116. t.Run(tt.name, func(t *testing.T) {
  1117. account := expiredGrokOAuthAccountForCredentialTest(int64(7500 + index))
  1118. baseRepo := &tokenRefreshAccountRepo{}
  1119. baseRepo.accountsByID = map[int64]*Account{account.ID: account}
  1120. repo := &grokCredentialCommitThenCancelRepo{
  1121. tokenRefreshAccountRepo: baseRepo,
  1122. returnErr: context.DeadlineExceeded,
  1123. }
  1124. cache := &grokTokenCacheForProviderTest{}
  1125. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, cache)}
  1126. token, err := svc.applyGrokCredentialAccountFailure(context.Background(), account, tt.class)
  1127. require.NoError(t, err, "the detached readback must resolve the inner timeout's commit ambiguity")
  1128. require.Empty(t, token)
  1129. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  1130. if tt.class.permanent {
  1131. require.Equal(t, StatusError, account.Status)
  1132. require.False(t, account.Schedulable)
  1133. require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys)
  1134. } else {
  1135. require.NotNil(t, account.TempUnschedulableUntil)
  1136. require.Equal(t, string(GrokCredentialReasonRefreshTransient), account.TempUnschedulableReason)
  1137. require.Empty(t, cache.deletedKeys)
  1138. }
  1139. })
  1140. }
  1141. }
  1142. func TestGrokCredentialUnconfirmedInnerStateDeadlineStopsAndRetainsSafetyBlock(t *testing.T) {
  1143. tests := []struct {
  1144. name string
  1145. class grokCredentialFailureClass
  1146. }{
  1147. {name: "permanent quarantine", class: grokCredentialFailureClass{reason: GrokCredentialReasonRevoked, permanent: true}},
  1148. {name: "temporary quarantine", class: grokCredentialFailureClass{reason: GrokCredentialReasonRefreshTransient, transient: true}},
  1149. }
  1150. for index, tt := range tests {
  1151. t.Run(tt.name, func(t *testing.T) {
  1152. account := expiredGrokOAuthAccountForCredentialTest(int64(7600 + index))
  1153. baseRepo := &tokenRefreshAccountRepo{}
  1154. baseRepo.accountsByID = map[int64]*Account{account.ID: account}
  1155. repo := &grokCredentialUncommittedDeadlineRepo{tokenRefreshAccountRepo: baseRepo}
  1156. cache := &grokTokenCacheForProviderTest{}
  1157. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, cache)}
  1158. token, err := svc.applyGrokCredentialAccountFailure(context.Background(), account, tt.class)
  1159. require.ErrorIs(t, err, errGrokCredentialStateUpdateFailed)
  1160. require.Empty(t, token)
  1161. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account), "an unknown commit outcome must retain the local safety block")
  1162. require.Equal(t, StatusActive, account.Status)
  1163. require.True(t, account.Schedulable)
  1164. require.Nil(t, account.TempUnschedulableUntil)
  1165. require.Empty(t, cache.deletedKeys)
  1166. })
  1167. }
  1168. }
  1169. func TestGrokCredentialRuntimeRollbackOwnership(t *testing.T) {
  1170. account := expiredGrokOAuthAccountForCredentialTest(734)
  1171. t.Run("later extending block survives", func(t *testing.T) {
  1172. svc := &OpenAIGatewayService{}
  1173. until := time.Now().Add(time.Minute)
  1174. rollbackFirst := svc.blockGrokCredentialRuntime(account, until, "first")
  1175. secondInstalled := make(chan struct{})
  1176. go func() {
  1177. svc.BlockAccountScheduling(account, until.Add(time.Minute), "independent")
  1178. close(secondInstalled)
  1179. }()
  1180. <-secondInstalled
  1181. rollbackFirst()
  1182. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account),
  1183. "rollback owned by the first invocation must not remove a later extending block")
  1184. })
  1185. t.Run("independent shorter block steals rollback ownership", func(t *testing.T) {
  1186. svc := &OpenAIGatewayService{}
  1187. until := time.Now().Add(2 * time.Minute)
  1188. rollbackFirst := svc.blockGrokCredentialRuntime(account, until, "first")
  1189. svc.BlockAccountScheduling(account, until.Add(-time.Minute), "shorter-no-op")
  1190. rollbackFirst()
  1191. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  1192. })
  1193. t.Run("serialized tentative rollbacks leave no block", func(t *testing.T) {
  1194. svc := &OpenAIGatewayService{}
  1195. for i := 0; i < 2; i++ {
  1196. mu := svc.grokCredentialMutationLock(account.ID)
  1197. require.NoError(t, mu.Lock(context.Background()))
  1198. rollback := svc.blockGrokCredentialRuntime(account, time.Now().Add(time.Minute), "tentative")
  1199. rollback()
  1200. mu.Unlock()
  1201. }
  1202. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  1203. })
  1204. }
  1205. func TestGetRequestCredentialAPIKeyBypassesOAuthFailureMapping(t *testing.T) {
  1206. gin.SetMode(gin.TestMode)
  1207. account := &Account{
  1208. ID: 705,
  1209. Platform: PlatformGrok,
  1210. Type: AccountTypeAPIKey,
  1211. Credentials: map[string]any{
  1212. "api_key": "third-party-key",
  1213. "base_url": "https://grok.example.test/v1",
  1214. },
  1215. }
  1216. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  1217. svc := &OpenAIGatewayService{}
  1218. token, kind, err := svc.getRequestCredential(context.Background(), c, account)
  1219. require.NoError(t, err)
  1220. require.Equal(t, "third-party-key", token)
  1221. require.Equal(t, "apikey", kind)
  1222. _, hasEvents := c.Get(OpsUpstreamErrorsKey)
  1223. require.False(t, hasEvents)
  1224. }
  1225. func TestPermanentCredentialFailureDoesNotDisableConcurrentlyRefreshedAccount(t *testing.T) {
  1226. account := expiredGrokOAuthAccountForCredentialTest(707)
  1227. latest := *account
  1228. latest.Credentials = shallowCopyMap(account.Credentials)
  1229. latest.Credentials["access_token"] = "fresh-access-token"
  1230. latest.Credentials["refresh_token"] = "rotated-refresh-token"
  1231. latest.Credentials["expires_at"] = time.Now().Add(time.Hour).UTC().Format(time.RFC3339)
  1232. latest.Credentials["_token_version"] = time.Now().UnixMilli()
  1233. repo := &tokenRefreshAccountRepo{}
  1234. repo.accountsByID = map[int64]*Account{account.ID: &latest}
  1235. cache := &grokTokenCacheForProviderTest{}
  1236. svc := &OpenAIGatewayService{
  1237. accountRepo: repo,
  1238. grokTokenProvider: NewGrokTokenProvider(repo, cache),
  1239. }
  1240. _, mutationErr := svc.applyGrokCredentialAccountFailure(context.Background(), account, grokCredentialFailureClass{
  1241. scope: GatewayFailureScopeAccount,
  1242. reason: GrokCredentialReasonRevoked,
  1243. action: NextAccountRetry,
  1244. permanent: true,
  1245. })
  1246. require.NoError(t, mutationErr)
  1247. require.Zero(t, repo.setErrorCalls)
  1248. require.Empty(t, cache.deletedKeys)
  1249. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  1250. }
  1251. func TestCredentialFailureConditionalMutationLosesToConcurrentRefresh(t *testing.T) {
  1252. for index, tt := range []struct {
  1253. name string
  1254. refreshErr error
  1255. }{
  1256. {name: "permanent", refreshErr: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_grant")},
  1257. {name: "transient", refreshErr: errors.New("temporary refresh transport failure")},
  1258. } {
  1259. t.Run(tt.name, func(t *testing.T) {
  1260. account := expiredGrokOAuthAccountForCredentialTest(int64(770 + index))
  1261. repo := &tokenRefreshAccountRepo{}
  1262. repo.accountsByID = map[int64]*Account{account.ID: account}
  1263. repo.beforeConditionalState = func() {
  1264. fresh := *account
  1265. fresh.Credentials = shallowCopyMap(account.Credentials)
  1266. fresh.Credentials["access_token"] = "refresh-won-token"
  1267. fresh.Credentials["refresh_token"] = "refresh-won-refresh"
  1268. fresh.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339)
  1269. fresh.Credentials["_token_version"] = time.Now().UnixMilli()
  1270. repo.accountsByID[account.ID] = &fresh
  1271. }
  1272. cache := &grokTokenCacheForProviderTest{lockResult: true}
  1273. provider := NewGrokTokenProvider(repo, cache)
  1274. provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{err: tt.refreshErr})
  1275. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  1276. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  1277. token, kind, err := svc.getRequestCredential(context.Background(), c, account)
  1278. require.NoError(t, err)
  1279. require.Equal(t, "refresh-won-token", token)
  1280. require.Equal(t, "oauth", kind)
  1281. require.Zero(t, repo.setErrorCalls)
  1282. require.Zero(t, repo.setTempUnschedCalls)
  1283. require.Empty(t, cache.deletedKeys)
  1284. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  1285. })
  1286. }
  1287. }
  1288. func TestCredentialFailureConditionalMutationLosesToConcurrentProxyRepair(t *testing.T) {
  1289. for index, tt := range []struct {
  1290. name string
  1291. refreshErr error
  1292. }{
  1293. {name: "permanent", refreshErr: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_grant")},
  1294. {name: "transient", refreshErr: errors.New("temporary refresh transport failure")},
  1295. } {
  1296. t.Run(tt.name, func(t *testing.T) {
  1297. account := expiredGrokOAuthAccountForCredentialTest(int64(780 + index))
  1298. oldProxyID := int64(10)
  1299. account.ProxyID = &oldProxyID
  1300. account.Proxy = &Proxy{}
  1301. repo := &tokenRefreshAccountRepo{}
  1302. repo.accountsByID = map[int64]*Account{account.ID: account}
  1303. repo.beforeConditionalState = func() {
  1304. fresh := *account
  1305. fresh.Credentials = shallowCopyMap(account.Credentials)
  1306. repairedProxyID := int64(11)
  1307. fresh.ProxyID = &repairedProxyID
  1308. repo.accountsByID[account.ID] = &fresh
  1309. }
  1310. cache := &grokTokenCacheForProviderTest{lockResult: true}
  1311. provider := NewGrokTokenProvider(repo, cache)
  1312. provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{err: tt.refreshErr})
  1313. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  1314. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  1315. _, _, err := svc.getRequestCredential(context.Background(), c, account)
  1316. var failoverErr *UpstreamFailoverError
  1317. require.ErrorAs(t, err, &failoverErr)
  1318. require.Equal(t, GrokCredentialReasonAccountChanged, failoverErr.Reason)
  1319. require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
  1320. require.Zero(t, repo.setErrorCalls)
  1321. require.Zero(t, repo.setTempUnschedCalls)
  1322. require.Empty(t, cache.deletedKeys)
  1323. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  1324. })
  1325. }
  1326. }
  1327. func TestCredentialFailureConditionalMutationLosesToSameIDProxyRestoration(t *testing.T) {
  1328. account := expiredGrokOAuthAccountForCredentialTest(790)
  1329. proxyID := int64(10)
  1330. account.ProxyID = &proxyID
  1331. account.Proxy = nil
  1332. repo := &tokenRefreshAccountRepo{}
  1333. repo.accountsByID = map[int64]*Account{account.ID: account}
  1334. repo.beforeConditionalState = func() {
  1335. account.Proxy = &Proxy{ID: proxyID}
  1336. }
  1337. cache := &grokTokenCacheForProviderTest{lockResult: true}
  1338. provider := NewGrokTokenProvider(repo, cache)
  1339. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  1340. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  1341. _, _, err := svc.getRequestCredential(context.Background(), c, account)
  1342. var failoverErr *UpstreamFailoverError
  1343. require.ErrorAs(t, err, &failoverErr)
  1344. require.Equal(t, GrokCredentialReasonAccountChanged, failoverErr.Reason)
  1345. require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
  1346. require.Zero(t, repo.setErrorCalls)
  1347. require.Zero(t, repo.setTempUnschedCalls)
  1348. require.Empty(t, cache.deletedKeys)
  1349. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  1350. }
  1351. func TestCredentialFailureConditionalMutationLosesToConcurrentUnschedulableState(t *testing.T) {
  1352. future := time.Now().Add(time.Hour)
  1353. states := []struct {
  1354. name string
  1355. mutate func(*Account)
  1356. }{
  1357. {name: "admin schedulable false", mutate: func(account *Account) { account.Schedulable = false }},
  1358. {name: "temporary cooldown", mutate: func(account *Account) { account.TempUnschedulableUntil = &future }},
  1359. {name: "rate limit cooldown", mutate: func(account *Account) { account.RateLimitResetAt = &future }},
  1360. {name: "overload cooldown", mutate: func(account *Account) { account.OverloadUntil = &future }},
  1361. }
  1362. classes := []struct {
  1363. name string
  1364. class grokCredentialFailureClass
  1365. }{
  1366. {name: "permanent", class: grokCredentialFailureClass{reason: GrokCredentialReasonRevoked, permanent: true}},
  1367. {name: "transient", class: grokCredentialFailureClass{reason: GrokCredentialReasonRefreshTransient, transient: true}},
  1368. }
  1369. for classIndex, classCase := range classes {
  1370. for stateIndex, stateCase := range states {
  1371. t.Run(classCase.name+"/"+stateCase.name, func(t *testing.T) {
  1372. account := expiredGrokOAuthAccountForCredentialTest(int64(791 + classIndex*10 + stateIndex))
  1373. repo := &tokenRefreshAccountRepo{}
  1374. repo.accountsByID = map[int64]*Account{account.ID: account}
  1375. repo.beforeConditionalState = func() { stateCase.mutate(account) }
  1376. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, &grokTokenCacheForProviderTest{})}
  1377. token, err := svc.applyGrokCredentialAccountFailure(context.Background(), account, classCase.class)
  1378. require.ErrorIs(t, err, errOAuthRefreshAccountStateChanged)
  1379. require.Empty(t, token)
  1380. require.Zero(t, repo.setErrorCalls)
  1381. require.Zero(t, repo.setTempUnschedCalls)
  1382. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  1383. })
  1384. }
  1385. }
  1386. }
  1387. func TestCredentialFailureCASMissDoesNotRecoverIneligibleLatestCredential(t *testing.T) {
  1388. future := time.Now().Add(time.Hour)
  1389. states := []struct {
  1390. name string
  1391. mutate func(*OpenAIGatewayService, *Account)
  1392. wantRuntimeBlocked bool
  1393. }{
  1394. {name: "disabled", mutate: func(_ *OpenAIGatewayService, account *Account) { account.Status = StatusDisabled }},
  1395. {name: "not schedulable", mutate: func(_ *OpenAIGatewayService, account *Account) { account.Schedulable = false }},
  1396. {name: "temporarily unschedulable", mutate: func(_ *OpenAIGatewayService, account *Account) { account.TempUnschedulableUntil = &future }},
  1397. {name: "rate limited", mutate: func(_ *OpenAIGatewayService, account *Account) { account.RateLimitResetAt = &future }},
  1398. {name: "overloaded", mutate: func(_ *OpenAIGatewayService, account *Account) { account.OverloadUntil = &future }},
  1399. {
  1400. name: "independently runtime blocked",
  1401. mutate: func(svc *OpenAIGatewayService, account *Account) {
  1402. svc.BlockAccountScheduling(account, time.Now().Add(24*time.Hour), "independent")
  1403. },
  1404. wantRuntimeBlocked: true,
  1405. },
  1406. }
  1407. classes := []struct {
  1408. name string
  1409. class grokCredentialFailureClass
  1410. }{
  1411. {name: "permanent", class: grokCredentialFailureClass{reason: GrokCredentialReasonRevoked, permanent: true}},
  1412. {name: "transient", class: grokCredentialFailureClass{reason: GrokCredentialReasonRefreshTransient, transient: true}},
  1413. }
  1414. for classIndex, classCase := range classes {
  1415. for stateIndex, stateCase := range states {
  1416. t.Run(classCase.name+"/"+stateCase.name, func(t *testing.T) {
  1417. account := expiredGrokOAuthAccountForCredentialTest(int64(800 + classIndex*20 + stateIndex))
  1418. repo := &tokenRefreshAccountRepo{}
  1419. repo.accountsByID = map[int64]*Account{account.ID: account}
  1420. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, &grokTokenCacheForProviderTest{})}
  1421. repo.beforeConditionalState = func() {
  1422. latest := *account
  1423. latest.Credentials = shallowCopyMap(account.Credentials)
  1424. latest.Credentials["access_token"] = "fresh-but-ineligible-token"
  1425. latest.Credentials["refresh_token"] = "fresh-but-ineligible-refresh"
  1426. latest.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339)
  1427. latest.Credentials["_token_version"] = time.Now().UnixMilli()
  1428. stateCase.mutate(svc, &latest)
  1429. repo.accountsByID[account.ID] = &latest
  1430. }
  1431. token, err := svc.applyGrokCredentialAccountFailure(context.Background(), account, classCase.class)
  1432. require.ErrorIs(t, err, errOAuthRefreshAccountStateChanged)
  1433. require.Empty(t, token)
  1434. require.Zero(t, repo.setErrorCalls)
  1435. require.Zero(t, repo.setTempUnschedCalls)
  1436. require.Equal(t, stateCase.wantRuntimeBlocked, svc.isOpenAIAccountRuntimeBlocked(account))
  1437. })
  1438. }
  1439. }
  1440. }
  1441. func TestGetRequestCredentialSharedCredentialPersistenceFailureStopsWithoutAccountMutation(t *testing.T) {
  1442. account := expiredGrokOAuthAccountForCredentialTest(782)
  1443. repo := &tokenRefreshAccountRepo{conditionalSuccessErr: errors.New("database unavailable")}
  1444. repo.accountsByID = map[int64]*Account{account.ID: account}
  1445. cache := &grokTokenCacheForProviderTest{lockResult: true}
  1446. provider := NewGrokTokenProvider(repo, cache)
  1447. provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{credentials: map[string]any{
  1448. "access_token": "new-access-token",
  1449. "refresh_token": "new-refresh-token",
  1450. "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
  1451. }})
  1452. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  1453. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  1454. _, _, err := svc.getRequestCredential(context.Background(), c, account)
  1455. var failoverErr *UpstreamFailoverError
  1456. require.ErrorAs(t, err, &failoverErr)
  1457. require.Equal(t, GatewayFailureScopeProvider, failoverErr.Scope)
  1458. require.Equal(t, GrokCredentialReasonProviderDown, failoverErr.Reason)
  1459. require.Equal(t, NextAccountStop, failoverErr.NextAccountAction)
  1460. require.Equal(t, 1, repo.conditionalSuccessCalls)
  1461. require.Zero(t, repo.updateCredentialsCalls)
  1462. require.Zero(t, repo.setErrorCalls)
  1463. require.Zero(t, repo.setTempUnschedCalls)
  1464. require.Empty(t, cache.deletedKeys)
  1465. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  1466. }
  1467. func TestGetRequestCredentialRecoversConcurrentRefreshWithoutFailover(t *testing.T) {
  1468. gin.SetMode(gin.TestMode)
  1469. account := expiredGrokOAuthAccountForCredentialTest(733)
  1470. latest := *account
  1471. latest.Credentials = shallowCopyMap(account.Credentials)
  1472. latest.Credentials["access_token"] = "fresh-concurrent-access"
  1473. latest.Credentials["refresh_token"] = "fresh-concurrent-refresh"
  1474. latest.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339)
  1475. latest.Credentials["_token_version"] = time.Now().UnixMilli()
  1476. baseRepo := &tokenRefreshAccountRepo{}
  1477. baseRepo.accountsByID = map[int64]*Account{account.ID: account}
  1478. repo := &grokCredentialSequencedRepo{tokenRefreshAccountRepo: baseRepo, latest: &latest}
  1479. cache := &grokTokenCacheForProviderTest{lockResult: true}
  1480. provider := NewGrokTokenProvider(repo, cache)
  1481. provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{
  1482. err: infraerrors.New(http.StatusForbidden, "GROK_OAUTH_ENTITLEMENT_DENIED", "access_denied"),
  1483. })
  1484. svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
  1485. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  1486. token, kind, err := svc.getRequestCredential(context.Background(), c, account)
  1487. require.NoError(t, err)
  1488. require.Equal(t, "fresh-concurrent-access", token)
  1489. require.Equal(t, "oauth", kind)
  1490. require.Zero(t, baseRepo.setErrorCalls)
  1491. require.Zero(t, baseRepo.setTempUnschedCalls)
  1492. require.Empty(t, cache.deletedKeys)
  1493. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  1494. _, hasEvents := c.Get(OpsUpstreamErrorsKey)
  1495. require.False(t, hasEvents)
  1496. }
  1497. func expiredGrokOAuthAccountForCredentialTest(id int64) *Account {
  1498. return &Account{
  1499. ID: id,
  1500. Platform: PlatformGrok,
  1501. Type: AccountTypeOAuth,
  1502. Status: StatusActive,
  1503. Schedulable: true,
  1504. Credentials: map[string]any{
  1505. "access_token": "expired-access-token",
  1506. "refresh_token": "refresh-token",
  1507. "expires_at": time.Now().Add(-time.Minute).UTC().Format(time.RFC3339),
  1508. "base_url": xai.DefaultCLIBaseURL,
  1509. },
  1510. }
  1511. }