| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665 |
- //go:build unit
- package service
- import (
- "context"
- "errors"
- "net/http"
- "net/http/httptest"
- "sync"
- "testing"
- "time"
- infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
- "github.com/Wei-Shaw/sub2api/internal/pkg/xai"
- "github.com/gin-gonic/gin"
- "github.com/stretchr/testify/require"
- )
- type grokCredentialPersistingRepo struct {
- *tokenRefreshAccountRepo
- }
- func TestClassifyGrokCredentialFailureBillingExhaustionIsTransient(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(9901)
- for _, message := range []string{
- "Grok OAuth refresh failed: spending limit reached",
- "included free usage exhausted",
- "credits exhausted",
- } {
- class := classifyGrokCredentialFailure(account, errors.New(message))
- require.Equal(t, GrokCredentialReasonRefreshTransient, class.reason, message)
- require.True(t, class.transient, message)
- require.False(t, class.permanent, message)
- }
- }
- func (r *grokCredentialPersistingRepo) SetError(ctx context.Context, id int64, message string) error {
- if err := ctx.Err(); err != nil {
- return err
- }
- if err := r.tokenRefreshAccountRepo.SetError(ctx, id, message); err != nil {
- return err
- }
- if account := r.accountsByID[id]; account != nil {
- account.Status = StatusError
- account.Schedulable = false
- account.ErrorMessage = message
- }
- return nil
- }
- type grokCredentialProxyRepoStub struct {
- ProxyRepository
- proxy *Proxy
- err error
- }
- func (r *grokCredentialProxyRepoStub) GetByID(context.Context, int64) (*Proxy, error) {
- return r.proxy, r.err
- }
- type grokCredentialBlockingRepo struct {
- *tokenRefreshAccountRepo
- setErrorStarted chan struct{}
- setTempStarted chan struct{}
- onceError sync.Once
- onceTemp sync.Once
- }
- type grokCredentialCommitThenCancelRepo struct {
- *tokenRefreshAccountRepo
- returnErr error
- }
- type grokCredentialUncommittedDeadlineRepo struct {
- *tokenRefreshAccountRepo
- }
- func (r *grokCredentialUncommittedDeadlineRepo) SetGrokCredentialErrorIfMatch(
- context.Context,
- int64,
- GrokCredentialMutationSnapshot,
- string,
- ) (bool, error) {
- return false, context.DeadlineExceeded
- }
- func (r *grokCredentialUncommittedDeadlineRepo) SetGrokCredentialTempUnschedulableIfMatch(
- context.Context,
- int64,
- GrokCredentialMutationSnapshot,
- time.Time,
- string,
- ) (bool, error) {
- return false, context.DeadlineExceeded
- }
- func (r *grokCredentialCommitThenCancelRepo) SetGrokCredentialErrorIfMatch(
- ctx context.Context,
- id int64,
- _ GrokCredentialMutationSnapshot,
- reason string,
- ) (bool, error) {
- account := r.accountsByID[id]
- account.Status = StatusError
- account.Schedulable = false
- account.ErrorMessage = reason
- if r.returnErr != nil {
- return false, r.returnErr
- }
- <-ctx.Done()
- return false, ctx.Err()
- }
- func (r *grokCredentialCommitThenCancelRepo) SetGrokCredentialTempUnschedulableIfMatch(
- ctx context.Context,
- id int64,
- _ GrokCredentialMutationSnapshot,
- until time.Time,
- reason string,
- ) (bool, error) {
- account := r.accountsByID[id]
- account.TempUnschedulableUntil = &until
- account.TempUnschedulableReason = reason
- if r.returnErr != nil {
- return false, r.returnErr
- }
- <-ctx.Done()
- return false, ctx.Err()
- }
- func (r *grokCredentialBlockingRepo) SetError(ctx context.Context, _ int64, _ string) error {
- r.onceError.Do(func() { close(r.setErrorStarted) })
- <-ctx.Done()
- return ctx.Err()
- }
- func (r *grokCredentialBlockingRepo) SetTempUnschedulable(ctx context.Context, _ int64, _ time.Time, _ string) error {
- r.onceTemp.Do(func() { close(r.setTempStarted) })
- <-ctx.Done()
- return ctx.Err()
- }
- func (r *grokCredentialBlockingRepo) SetGrokCredentialErrorIfMatch(
- ctx context.Context,
- _ int64,
- _ GrokCredentialMutationSnapshot,
- _ string,
- ) (bool, error) {
- r.onceError.Do(func() { close(r.setErrorStarted) })
- <-ctx.Done()
- return false, ctx.Err()
- }
- func (r *grokCredentialBlockingRepo) SetGrokCredentialTempUnschedulableIfMatch(
- ctx context.Context,
- _ int64,
- _ GrokCredentialMutationSnapshot,
- _ time.Time,
- _ string,
- ) (bool, error) {
- r.onceTemp.Do(func() { close(r.setTempStarted) })
- <-ctx.Done()
- return false, ctx.Err()
- }
- type grokCredentialBlockingCache struct {
- GrokTokenCache
- deleteStarted chan struct{}
- releaseDelete chan struct{}
- once sync.Once
- mu sync.Mutex
- deleted bool
- }
- type grokCredentialSequencedRepo struct {
- *tokenRefreshAccountRepo
- mu sync.Mutex
- latest *Account
- getCall int
- }
- type grokCredentialRereadFailureRepo struct {
- *tokenRefreshAccountRepo
- account *Account
- err error
- }
- type grokCredentialCountingRefresher struct {
- refreshCalls int
- }
- func (r *grokCredentialCountingRefresher) CacheKey(account *Account) string {
- return GrokTokenCacheKey(account)
- }
- func (r *grokCredentialCountingRefresher) CanRefresh(*Account) bool { return true }
- func (r *grokCredentialCountingRefresher) NeedsRefresh(*Account, time.Duration) bool { return true }
- func (r *grokCredentialCountingRefresher) Refresh(context.Context, *Account) (map[string]any, error) {
- r.refreshCalls++
- return map[string]any{"access_token": "must-not-be-used"}, nil
- }
- func (r *grokCredentialRereadFailureRepo) GetByID(context.Context, int64) (*Account, error) {
- return r.account, r.err
- }
- func (r *grokCredentialSequencedRepo) GetByID(ctx context.Context, id int64) (*Account, error) {
- r.mu.Lock()
- defer r.mu.Unlock()
- r.getCall++
- if r.getCall > 1 && r.latest != nil {
- return r.latest, nil
- }
- return r.tokenRefreshAccountRepo.GetByID(ctx, id)
- }
- func (c *grokCredentialBlockingCache) DeleteAccessToken(ctx context.Context, _ string) error {
- c.once.Do(func() { close(c.deleteStarted) })
- select {
- case <-c.releaseDelete:
- c.mu.Lock()
- c.deleted = true
- c.mu.Unlock()
- return nil
- case <-ctx.Done():
- return ctx.Err()
- }
- }
- func (c *grokCredentialBlockingCache) wasDeleted() bool {
- c.mu.Lock()
- defer c.mu.Unlock()
- return c.deleted
- }
- func TestUpstreamFailoverErrorNextAccountActionPreservesLegacyRetry(t *testing.T) {
- t.Parallel()
- require.True(t, (&UpstreamFailoverError{}).ShouldRetryNextAccount())
- require.True(t, (&UpstreamFailoverError{NextAccountAction: NextAccountRetry}).ShouldRetryNextAccount())
- require.False(t, (&UpstreamFailoverError{NextAccountAction: NextAccountStop}).ShouldRetryNextAccount())
- }
- func TestGetRequestCredentialMapsPermanentGrokOAuthFailureAndRedactsSecrets(t *testing.T) {
- gin.SetMode(gin.TestMode)
- account := expiredGrokOAuthAccountForCredentialTest(701)
- repo := &tokenRefreshAccountRepo{}
- repo.accountsByID = map[int64]*Account{account.ID: account}
- cache := &grokTokenCacheForProviderTest{lockResult: true}
- provider := NewGrokTokenProvider(repo, cache)
- provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{
- err: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_grant access_token=leaked-access refresh_token=leaked-refresh"),
- })
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- token, kind, err := svc.getRequestCredential(context.Background(), c, account)
- require.Error(t, err)
- require.Empty(t, token)
- require.Empty(t, kind)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, GatewayFailureStageAccountAuth, failoverErr.Stage)
- require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope)
- require.Equal(t, GrokCredentialReasonRevoked, failoverErr.Reason)
- require.True(t, failoverErr.ShouldRetryNextAccount())
- require.Equal(t, 0, failoverErr.StatusCode)
- require.Equal(t, http.StatusServiceUnavailable, failoverErr.ClientStatusCode)
- require.NotContains(t, err.Error(), "leaked-access")
- require.NotContains(t, err.Error(), "leaked-refresh")
- require.Equal(t, 1, repo.setErrorCalls)
- require.Zero(t, repo.setTempUnschedCalls)
- require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
- require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys)
- require.NotContains(t, repo.lastErrorMessage, "leaked-access")
- require.NotContains(t, repo.lastErrorMessage, "leaked-refresh")
- rawEvents, ok := c.Get(OpsUpstreamErrorsKey)
- require.True(t, ok)
- events, ok := rawEvents.([]*OpsUpstreamErrorEvent)
- require.True(t, ok)
- require.Len(t, events, 1)
- require.Equal(t, string(GatewayFailureStageAccountAuth), events[0].Stage)
- require.Equal(t, string(GatewayFailureScopeAccount), events[0].Scope)
- require.Equal(t, string(GrokCredentialReasonRevoked), events[0].Reason)
- require.Zero(t, events[0].UpstreamStatusCode)
- require.NotContains(t, events[0].Message, "leaked-access")
- require.NotContains(t, events[0].Message, "leaked-refresh")
- }
- func TestNewGrokCredentialFailoverDoesNotAttributeInferenceProxy(t *testing.T) {
- gin.SetMode(gin.TestMode)
- proxyID := int64(43)
- account := &Account{
- ID: 701,
- Platform: PlatformGrok,
- ProxyID: &proxyID,
- Proxy: &Proxy{ID: proxyID, Name: "bound-proxy"},
- }
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- err := (&OpenAIGatewayService{}).newGrokCredentialFailover(c, account, grokCredentialFailureClass{
- scope: GatewayFailureScopeAccount,
- reason: GrokCredentialReasonAccountChanged,
- action: NextAccountRetry,
- message: "credential unavailable",
- })
- require.Error(t, err)
- rawEvents, ok := c.Get(OpsUpstreamErrorsKey)
- require.True(t, ok)
- events, ok := rawEvents.([]*OpsUpstreamErrorEvent)
- require.True(t, ok)
- require.Len(t, events, 1)
- require.Nil(t, events[0].ProxyID)
- require.Equal(t, opsProxyNameUnknown, events[0].ProxyName)
- }
- func TestGetRequestCredentialPermanentMappingsPersistAndInvalidate(t *testing.T) {
- gin.SetMode(gin.TestMode)
- tests := []struct {
- name string
- prepare func(*Account)
- refreshErr error
- wantReason GatewayFailureReason
- cachedToken string
- }{
- {
- name: "missing refresh credential",
- prepare: func(account *Account) {
- delete(account.Credentials, "refresh_token")
- },
- wantReason: GrokCredentialReasonMissing,
- },
- {
- name: "missing access credential",
- prepare: func(account *Account) {
- delete(account.Credentials, "access_token")
- account.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339)
- },
- wantReason: GrokCredentialReasonMissing,
- cachedToken: "stale-cached-access",
- },
- {
- name: "explicit entitlement action required",
- prepare: func(*Account) {},
- refreshErr: infraerrors.New(http.StatusForbidden, "GROK_OAUTH_ENTITLEMENT_DENIED", "access_denied"),
- wantReason: GrokCredentialReasonEntitlement,
- },
- }
- for index, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(int64(720 + index))
- account.Status = StatusActive
- account.Schedulable = true
- tt.prepare(account)
- baseRepo := &tokenRefreshAccountRepo{}
- baseRepo.accountsByID = map[int64]*Account{account.ID: account}
- repo := &grokCredentialPersistingRepo{tokenRefreshAccountRepo: baseRepo}
- cache := &grokTokenCacheForProviderTest{lockResult: true, token: tt.cachedToken}
- provider := NewGrokTokenProvider(repo, cache)
- provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{err: tt.refreshErr})
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- _, _, err := svc.getRequestCredential(context.Background(), c, account)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, tt.wantReason, failoverErr.Reason)
- require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope)
- require.Equal(t, 1, baseRepo.setErrorCalls)
- require.Equal(t, StatusError, account.Status)
- require.False(t, account.Schedulable)
- require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
- require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys)
- })
- }
- }
- func TestGetRequestCredentialMissingAccessNeverRefreshesAndPermanentlyFailsOver(t *testing.T) {
- gin.SetMode(gin.TestMode)
- tests := []struct {
- name string
- expiresAt *time.Time
- }{
- {name: "expiry missing"},
- {name: "expired", expiresAt: func() *time.Time { value := time.Now().Add(-time.Minute); return &value }()},
- {name: "near expiry", expiresAt: func() *time.Time { value := time.Now().Add(30 * time.Minute); return &value }()},
- }
- for index, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(int64(760 + index))
- account.Schedulable = true
- delete(account.Credentials, "access_token")
- if tt.expiresAt == nil {
- delete(account.Credentials, "expires_at")
- } else {
- account.Credentials["expires_at"] = tt.expiresAt.UTC().Format(time.RFC3339)
- }
- baseRepo := &tokenRefreshAccountRepo{}
- baseRepo.accountsByID = map[int64]*Account{account.ID: account}
- repo := &grokCredentialPersistingRepo{tokenRefreshAccountRepo: baseRepo}
- cache := &grokTokenCacheForProviderTest{lockResult: true, token: "stale-cache-must-not-win"}
- refresher := &grokCredentialCountingRefresher{}
- provider := NewGrokTokenProvider(repo, cache)
- provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), refresher)
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- token, kind, err := svc.getRequestCredential(context.Background(), c, account)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Empty(t, token)
- require.Empty(t, kind)
- require.Equal(t, GrokCredentialReasonMissing, failoverErr.Reason)
- require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
- require.Zero(t, refresher.refreshCalls, "structurally missing access credentials must not reach the token endpoint")
- require.Equal(t, 1, baseRepo.setErrorCalls)
- require.Equal(t, StatusError, account.Status)
- require.False(t, account.Schedulable)
- require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys)
- require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
- })
- }
- }
- func TestGetRequestCredentialWarmCachedAccessWithMissingRefreshPermanentlyFailsOver(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(764)
- account.Credentials["access_token"] = "valid-access"
- account.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339)
- delete(account.Credentials, "refresh_token")
- baseRepo := &tokenRefreshAccountRepo{}
- baseRepo.accountsByID = map[int64]*Account{account.ID: account}
- repo := &grokCredentialPersistingRepo{tokenRefreshAccountRepo: baseRepo}
- cache := &grokTokenCacheForProviderTest{lockResult: true, token: "valid-access"}
- refresher := &grokCredentialCountingRefresher{}
- provider := NewGrokTokenProvider(repo, cache)
- provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), refresher)
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- _, _, err := svc.getRequestCredential(context.Background(), c, account)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, GrokCredentialReasonMissing, failoverErr.Reason)
- require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
- require.Zero(t, refresher.refreshCalls)
- require.Equal(t, 1, baseRepo.setErrorCalls)
- require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys)
- require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
- }
- func TestGetRequestCredentialMapsTransientAndProviderFailuresSeparately(t *testing.T) {
- gin.SetMode(gin.TestMode)
- t.Run("account transient temporarily unschedules", func(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(702)
- repo := &tokenRefreshAccountRepo{}
- repo.accountsByID = map[int64]*Account{account.ID: account}
- cache := &grokTokenCacheForProviderTest{lockResult: true}
- provider := NewGrokTokenProvider(repo, cache)
- provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{err: errors.New("temporary refresh transport failure")})
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- _, _, err := svc.getRequestCredential(context.Background(), c, account)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope)
- require.Equal(t, GrokCredentialReasonRefreshTransient, failoverErr.Reason)
- require.True(t, failoverErr.ShouldRetryNextAccount())
- require.Zero(t, repo.setErrorCalls)
- require.Equal(t, 1, repo.setTempUnschedCalls)
- require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
- })
- t.Run("shared provider configuration stops without mutation", func(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(703)
- repo := &tokenRefreshAccountRepo{}
- repo.accountsByID = map[int64]*Account{account.ID: account}
- provider := NewGrokTokenProvider(repo, nil)
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- _, _, err := svc.getRequestCredential(context.Background(), c, account)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, GatewayFailureScopeProvider, failoverErr.Scope)
- require.Equal(t, NextAccountStop, failoverErr.NextAccountAction)
- require.Zero(t, repo.setErrorCalls)
- require.Zero(t, repo.setTempUnschedCalls)
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- })
- t.Run("account reread failures preserve shared versus missing-row scope", func(t *testing.T) {
- for _, tt := range []struct {
- name string
- account *Account
- err error
- }{
- {name: "repository error", err: errors.New("database temporarily unavailable")},
- {name: "missing row"},
- } {
- t.Run(tt.name, func(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(712)
- baseRepo := &tokenRefreshAccountRepo{}
- baseRepo.accountsByID = map[int64]*Account{account.ID: account}
- repo := &grokCredentialRereadFailureRepo{tokenRefreshAccountRepo: baseRepo, account: tt.account, err: tt.err}
- cache := &grokTokenCacheForProviderTest{lockResult: true}
- provider := NewGrokTokenProvider(repo, cache)
- provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), NewGrokTokenRefresher(nil))
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- _, _, err := svc.getRequestCredential(context.Background(), c, account)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- if tt.err != nil {
- require.Equal(t, GatewayFailureScopeProvider, failoverErr.Scope)
- require.Equal(t, GrokCredentialReasonProviderDown, failoverErr.Reason)
- require.Equal(t, NextAccountStop, failoverErr.NextAccountAction)
- } else {
- require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope)
- require.Equal(t, GrokCredentialReasonAccountChanged, failoverErr.Reason)
- require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
- }
- require.Zero(t, baseRepo.setErrorCalls)
- require.Zero(t, baseRepo.setTempUnschedCalls)
- require.Empty(t, cache.deletedKeys)
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- })
- }
- })
- t.Run("fresh account eligibility changes retry without mutating stale state", func(t *testing.T) {
- for _, tt := range []struct {
- name string
- mutate func(*Account)
- }{
- {
- name: "account disabled",
- mutate: func(account *Account) {
- account.Status = StatusDisabled
- },
- },
- {
- name: "account converted",
- mutate: func(account *Account) {
- account.Type = AccountTypeUpstream
- },
- },
- {
- name: "account manually unschedulable",
- mutate: func(account *Account) {
- account.Schedulable = false
- },
- },
- {
- name: "account temporarily unschedulable",
- mutate: func(account *Account) {
- until := time.Now().Add(time.Minute)
- account.TempUnschedulableUntil = &until
- },
- },
- } {
- t.Run(tt.name, func(t *testing.T) {
- staleAccount := expiredGrokOAuthAccountForCredentialTest(713)
- freshAccount := *staleAccount
- freshAccount.Credentials = shallowCopyMap(staleAccount.Credentials)
- tt.mutate(&freshAccount)
- repo := &tokenRefreshAccountRepo{}
- repo.accountsByID = map[int64]*Account{staleAccount.ID: &freshAccount}
- cache := &grokTokenCacheForProviderTest{lockResult: true}
- provider := NewGrokTokenProvider(repo, cache)
- provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), NewGrokTokenRefresher(nil))
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- _, _, err := svc.getRequestCredential(context.Background(), c, staleAccount)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope)
- require.Equal(t, GrokCredentialReasonAccountChanged, failoverErr.Reason)
- require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
- require.Zero(t, repo.setErrorCalls)
- require.Zero(t, repo.setTempUnschedCalls)
- require.Empty(t, cache.deletedKeys)
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(staleAccount))
- })
- }
- })
- t.Run("fresh missing refresh credential permanently blocks the account", func(t *testing.T) {
- staleAccount := expiredGrokOAuthAccountForCredentialTest(714)
- staleAccount.Schedulable = true
- freshAccount := *staleAccount
- freshAccount.Credentials = shallowCopyMap(staleAccount.Credentials)
- delete(freshAccount.Credentials, "refresh_token")
- baseRepo := &tokenRefreshAccountRepo{}
- baseRepo.accountsByID = map[int64]*Account{staleAccount.ID: &freshAccount}
- repo := &grokCredentialPersistingRepo{tokenRefreshAccountRepo: baseRepo}
- cache := &grokTokenCacheForProviderTest{lockResult: true}
- provider := NewGrokTokenProvider(repo, cache)
- provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), NewGrokTokenRefresher(nil))
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- _, _, err := svc.getRequestCredential(context.Background(), c, staleAccount)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope)
- require.Equal(t, GrokCredentialReasonMissing, failoverErr.Reason)
- require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
- require.Equal(t, 1, baseRepo.setErrorCalls)
- require.Zero(t, baseRepo.setTempUnschedCalls)
- require.Equal(t, StatusError, freshAccount.Status)
- require.False(t, freshAccount.Schedulable)
- require.Equal(t, []string{GrokTokenCacheKey(staleAccount)}, cache.deletedKeys)
- require.True(t, svc.isOpenAIAccountRuntimeBlocked(staleAccount))
- })
- t.Run("refresh added after locked structural failure wins conditional mutation", func(t *testing.T) {
- staleAccount := expiredGrokOAuthAccountForCredentialTest(717)
- freshAccount := *staleAccount
- freshAccount.Credentials = shallowCopyMap(staleAccount.Credentials)
- delete(freshAccount.Credentials, "refresh_token")
- repo := &tokenRefreshAccountRepo{}
- repo.accountsByID = map[int64]*Account{staleAccount.ID: &freshAccount}
- repo.beforeConditionalState = func() {
- repaired := freshAccount
- repaired.Credentials = shallowCopyMap(freshAccount.Credentials)
- repaired.Credentials["refresh_token"] = "repaired-refresh-token"
- repaired.Credentials["expires_at"] = time.Now().Add(time.Hour).UTC().Format(time.RFC3339)
- repo.accountsByID[staleAccount.ID] = &repaired
- }
- cache := &grokTokenCacheForProviderTest{lockResult: true}
- provider := NewGrokTokenProvider(repo, cache)
- provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), NewGrokTokenRefresher(nil))
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- token, kind, err := svc.getRequestCredential(context.Background(), c, staleAccount)
- require.NoError(t, err)
- require.Equal(t, "expired-access-token", token)
- require.Equal(t, "oauth", kind)
- require.Zero(t, repo.setErrorCalls)
- require.Zero(t, repo.setTempUnschedCalls)
- require.Empty(t, cache.deletedKeys)
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(staleAccount))
- })
- t.Run("expiry-only repair wins full credential fingerprint CAS", func(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(718)
- repo := &tokenRefreshAccountRepo{}
- repo.accountsByID = map[int64]*Account{account.ID: account}
- repo.beforeConditionalState = func() {
- repaired := *account
- repaired.Credentials = shallowCopyMap(account.Credentials)
- repaired.Credentials["expires_at"] = time.Now().Add(time.Hour).UTC().Format(time.RFC3339)
- repo.accountsByID[account.ID] = &repaired
- }
- cache := &grokTokenCacheForProviderTest{lockResult: true}
- provider := NewGrokTokenProvider(repo, cache)
- provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{
- err: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_grant"),
- })
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- token, kind, err := svc.getRequestCredential(context.Background(), c, account)
- require.NoError(t, err)
- require.Equal(t, "expired-access-token", token)
- require.Equal(t, "oauth", kind)
- require.Zero(t, repo.setErrorCalls)
- require.Empty(t, cache.deletedKeys)
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- })
- t.Run("generic token endpoint 403 stops as shared provider failure", func(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(708)
- repo := &tokenRefreshAccountRepo{}
- repo.accountsByID = map[int64]*Account{account.ID: account}
- cache := &grokTokenCacheForProviderTest{lockResult: true}
- provider := NewGrokTokenProvider(repo, cache)
- provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{
- err: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "token refresh failed: status 403, body: forbidden"),
- })
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- _, _, err := svc.getRequestCredential(context.Background(), c, account)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, GatewayFailureScopeProvider, failoverErr.Scope)
- require.Equal(t, GrokCredentialReasonProviderDown, failoverErr.Reason)
- require.Equal(t, NextAccountStop, failoverErr.NextAccountAction)
- require.Zero(t, repo.setErrorCalls)
- require.Zero(t, repo.setTempUnschedCalls)
- require.Empty(t, cache.deletedKeys)
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- })
- t.Run("account proxy generic 403 remains bounded account transient", func(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(711)
- proxyID := int64(43)
- account.ProxyID = &proxyID
- account.Proxy = &Proxy{}
- repo := &tokenRefreshAccountRepo{}
- repo.accountsByID = map[int64]*Account{account.ID: account}
- cache := &grokTokenCacheForProviderTest{lockResult: true}
- provider := NewGrokTokenProvider(repo, cache)
- provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{
- err: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "token refresh failed: status 403, body: forbidden"),
- })
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- _, _, err := svc.getRequestCredential(context.Background(), c, account)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope)
- require.Equal(t, GrokCredentialReasonRefreshTransient, failoverErr.Reason)
- require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
- require.Zero(t, repo.setErrorCalls)
- require.Equal(t, 1, repo.setTempUnschedCalls)
- require.Empty(t, cache.deletedKeys)
- require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
- })
- t.Run("proxy repository read failure stops without account mutation", func(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(709)
- proxyID := int64(41)
- account.ProxyID = &proxyID
- account.Proxy = &Proxy{}
- repo := &tokenRefreshAccountRepo{}
- repo.accountsByID = map[int64]*Account{account.ID: account}
- cache := &grokTokenCacheForProviderTest{lockResult: true}
- oauthSvc := NewGrokOAuthService(&grokCredentialProxyRepoStub{err: errors.New("database temporarily unavailable")}, &grokOAuthClientStub{})
- defer oauthSvc.Stop()
- provider := NewGrokTokenProvider(repo, cache)
- provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), NewGrokTokenRefresher(oauthSvc))
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- _, _, err := svc.getRequestCredential(context.Background(), c, account)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, GatewayFailureScopeProvider, failoverErr.Scope)
- require.Equal(t, GrokCredentialReasonProviderDown, failoverErr.Reason)
- require.Equal(t, NextAccountStop, failoverErr.NextAccountAction)
- require.Zero(t, repo.setErrorCalls)
- require.Zero(t, repo.setTempUnschedCalls)
- require.Empty(t, cache.deletedKeys)
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- })
- t.Run("structurally missing configured proxy permanently blocks only that account", func(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(710)
- account.Status = StatusActive
- account.Schedulable = true
- proxyID := int64(42)
- account.ProxyID = &proxyID
- baseRepo := &tokenRefreshAccountRepo{}
- baseRepo.accountsByID = map[int64]*Account{account.ID: account}
- repo := &grokCredentialPersistingRepo{tokenRefreshAccountRepo: baseRepo}
- cache := &grokTokenCacheForProviderTest{lockResult: true}
- oauthSvc := NewGrokOAuthService(&grokCredentialProxyRepoStub{err: ErrProxyNotFound}, &grokOAuthClientStub{})
- defer oauthSvc.Stop()
- provider := NewGrokTokenProvider(repo, cache)
- provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), NewGrokTokenRefresher(oauthSvc))
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- _, _, err := svc.getRequestCredential(context.Background(), c, account)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope)
- require.Equal(t, GrokCredentialReasonProxyInvalid, failoverErr.Reason)
- require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
- require.Equal(t, 1, baseRepo.setErrorCalls)
- require.Equal(t, StatusError, account.Status)
- require.False(t, account.Schedulable)
- require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys)
- require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
- })
- }
- func TestGetRequestCredentialRuntimeBlockWinsBeforeWarmTokenCache(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(716)
- account.Credentials["access_token"] = "valid-access"
- account.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339)
- repo := &tokenRefreshAccountRepo{}
- repo.accountsByID = map[int64]*Account{account.ID: account}
- cache := &grokTokenCacheForProviderTest{token: "valid-access"}
- provider := NewGrokTokenProvider(repo, cache)
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- svc.BlockAccountScheduling(account, time.Now().Add(time.Minute), "independent")
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- _, _, err := svc.getRequestCredential(context.Background(), c, account)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, GrokCredentialReasonAccountChanged, failoverErr.Reason)
- require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
- require.Zero(t, cache.getCalls)
- require.Zero(t, repo.setErrorCalls)
- require.Zero(t, repo.setTempUnschedCalls)
- }
- func TestGetRequestCredentialWarmCachedAccessWithMissingConfiguredProxyPermanentlyFailsOver(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(715)
- account.Credentials["access_token"] = "valid-access"
- account.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339)
- proxyID := int64(44)
- account.ProxyID = &proxyID
- account.Proxy = nil
- baseRepo := &tokenRefreshAccountRepo{}
- baseRepo.accountsByID = map[int64]*Account{account.ID: account}
- repo := &grokCredentialPersistingRepo{tokenRefreshAccountRepo: baseRepo}
- cache := &grokTokenCacheForProviderTest{lockResult: true, token: "valid-access"}
- refresher := &grokCredentialCountingRefresher{}
- provider := NewGrokTokenProvider(repo, cache)
- provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), refresher)
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- _, _, err := svc.getRequestCredential(context.Background(), c, account)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, GrokCredentialReasonProxyInvalid, failoverErr.Reason)
- require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
- require.Zero(t, refresher.refreshCalls)
- require.Equal(t, 1, baseRepo.setErrorCalls)
- require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys)
- require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
- }
- func TestGetRequestCredentialCancellationAndBudgetDoNotMutateAccount(t *testing.T) {
- gin.SetMode(gin.TestMode)
- account := expiredGrokOAuthAccountForCredentialTest(704)
- repo := &tokenRefreshAccountRepo{}
- repo.accountsByID = map[int64]*Account{account.ID: account}
- provider := NewGrokTokenProvider(repo, nil)
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- t.Run("parent cancellation is returned directly", func(t *testing.T) {
- ctx, cancel := context.WithCancel(context.Background())
- cancel()
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- _, _, err := svc.getRequestCredential(ctx, c, account)
- require.ErrorIs(t, err, context.Canceled)
- var failoverErr *UpstreamFailoverError
- require.False(t, errors.As(err, &failoverErr))
- })
- t.Run("request credential budget stops safely", func(t *testing.T) {
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- c.Set(grokCredentialFailoverDeadlineKey, time.Now().Add(-time.Second))
- _, _, err := svc.getRequestCredential(context.Background(), c, account)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, GatewayFailureScopeRequest, failoverErr.Scope)
- require.Equal(t, GrokCredentialReasonFailoverTimeout, failoverErr.Reason)
- require.False(t, failoverErr.ShouldRetryNextAccount())
- })
- require.Zero(t, repo.setErrorCalls)
- require.Zero(t, repo.setTempUnschedCalls)
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- }
- func TestGetRequestCredentialStateMutationFailureStopsAndKeepsRuntimeBlock(t *testing.T) {
- gin.SetMode(gin.TestMode)
- tests := []struct {
- name string
- refreshErr error
- configure func(*tokenRefreshAccountRepo, *grokTokenCacheForProviderTest)
- wantSetError int
- wantSetTemp int
- wantCacheDelete int
- }{
- {
- name: "permanent state persistence",
- refreshErr: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_grant"),
- configure: func(repo *tokenRefreshAccountRepo, _ *grokTokenCacheForProviderTest) {
- repo.setErrorErr = errors.New("database write failed")
- },
- wantSetError: 1,
- },
- {
- name: "transient state persistence",
- refreshErr: errors.New("temporary refresh transport failure"),
- configure: func(repo *tokenRefreshAccountRepo, _ *grokTokenCacheForProviderTest) {
- repo.setTempUnschedErr = errors.New("database write failed")
- },
- wantSetTemp: 1,
- },
- {
- name: "permanent token cache invalidation",
- refreshErr: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_grant"),
- configure: func(_ *tokenRefreshAccountRepo, cache *grokTokenCacheForProviderTest) {
- cache.deleteErr = errors.New("cache delete failed")
- },
- wantSetError: 1,
- wantCacheDelete: 1,
- },
- }
- for index, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(int64(740 + index))
- repo := &tokenRefreshAccountRepo{}
- repo.accountsByID = map[int64]*Account{account.ID: account}
- cache := &grokTokenCacheForProviderTest{lockResult: true}
- tt.configure(repo, cache)
- provider := NewGrokTokenProvider(repo, cache)
- provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{err: tt.refreshErr})
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- _, _, err := svc.getRequestCredential(context.Background(), c, account)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, GatewayFailureScopeProvider, failoverErr.Scope)
- require.Equal(t, GrokCredentialReasonStateUpdate, failoverErr.Reason)
- require.Equal(t, NextAccountStop, failoverErr.NextAccountAction)
- require.Equal(t, tt.wantSetError, repo.setErrorCalls)
- require.Equal(t, tt.wantSetTemp, repo.setTempUnschedCalls)
- require.Len(t, cache.deletedKeys, tt.wantCacheDelete)
- require.True(t, svc.isOpenAIAccountRuntimeBlocked(account), "failed mutation must retain the immediate local block")
- })
- }
- }
- func TestGrokCredentialMutationBoundariesHonorParentCancellation(t *testing.T) {
- t.Run("blocked SetError cancellation prevents cache and runtime mutation", func(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(730)
- baseRepo := &tokenRefreshAccountRepo{}
- baseRepo.accountsByID = map[int64]*Account{account.ID: account}
- repo := &grokCredentialBlockingRepo{
- tokenRefreshAccountRepo: baseRepo,
- setErrorStarted: make(chan struct{}),
- setTempStarted: make(chan struct{}),
- }
- cache := &grokTokenCacheForProviderTest{}
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, cache)}
- ctx, cancel := context.WithCancel(context.Background())
- result := make(chan error, 1)
- go func() {
- _, err := svc.applyGrokCredentialAccountFailure(ctx, account, grokCredentialFailureClass{
- reason: GrokCredentialReasonRevoked, permanent: true,
- })
- result <- err
- }()
- <-repo.setErrorStarted
- require.True(t, svc.isOpenAIAccountRuntimeBlocked(account), "runtime block must precede persistent SetError")
- cancel()
- require.ErrorIs(t, <-result, context.Canceled)
- require.Zero(t, baseRepo.setErrorCalls)
- require.Empty(t, cache.deletedKeys)
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- })
- t.Run("blocked temporary unschedule cancellation prevents runtime mutation", func(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(731)
- baseRepo := &tokenRefreshAccountRepo{}
- baseRepo.accountsByID = map[int64]*Account{account.ID: account}
- repo := &grokCredentialBlockingRepo{
- tokenRefreshAccountRepo: baseRepo,
- setErrorStarted: make(chan struct{}),
- setTempStarted: make(chan struct{}),
- }
- svc := &OpenAIGatewayService{accountRepo: repo}
- ctx, cancel := context.WithCancel(context.Background())
- result := make(chan error, 1)
- go func() {
- _, err := svc.applyGrokCredentialAccountFailure(ctx, account, grokCredentialFailureClass{
- reason: GrokCredentialReasonRefreshTransient, transient: true,
- })
- result <- err
- }()
- <-repo.setTempStarted
- require.True(t, svc.isOpenAIAccountRuntimeBlocked(account), "runtime block must precede temporary unscheduling")
- cancel()
- require.ErrorIs(t, <-result, context.Canceled)
- require.Zero(t, baseRepo.setTempUnschedCalls)
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- })
- t.Run("post-commit cancellation finishes cache cleanup and retains quarantine", func(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(732)
- repo := &tokenRefreshAccountRepo{}
- repo.accountsByID = map[int64]*Account{account.ID: account}
- cache := &grokCredentialBlockingCache{deleteStarted: make(chan struct{}), releaseDelete: make(chan struct{})}
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, cache)}
- ctx, cancel := context.WithCancel(context.Background())
- result := make(chan error, 1)
- go func() {
- _, err := svc.applyGrokCredentialAccountFailure(ctx, account, grokCredentialFailureClass{
- reason: GrokCredentialReasonRevoked, permanent: true,
- })
- result <- err
- }()
- <-cache.deleteStarted
- require.True(t, svc.isOpenAIAccountRuntimeBlocked(account), "runtime block must precede cache invalidation")
- cancel()
- close(cache.releaseDelete)
- require.ErrorIs(t, <-result, context.Canceled)
- require.Equal(t, 1, repo.setErrorCalls)
- require.True(t, cache.wasDeleted())
- require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
- })
- }
- func TestGrokCredentialMutationLockWaitHonorsCredentialBudget(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(735)
- repo := &tokenRefreshAccountRepo{}
- repo.accountsByID = map[int64]*Account{account.ID: account}
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, &grokTokenCacheForProviderTest{})}
- mutationLock := svc.grokCredentialMutationLock(account.ID)
- require.NoError(t, mutationLock.Lock(context.Background()))
- defer mutationLock.Unlock()
- ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
- defer cancel()
- startedAt := time.Now()
- token, err := svc.applyGrokCredentialAccountFailure(ctx, account, grokCredentialFailureClass{
- reason: GrokCredentialReasonRevoked, permanent: true,
- })
- require.ErrorIs(t, err, context.DeadlineExceeded)
- require.Empty(t, token)
- require.Less(t, time.Since(startedAt), 500*time.Millisecond)
- require.Zero(t, repo.setErrorCalls)
- require.Zero(t, repo.setTempUnschedCalls)
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- }
- func TestGetRequestCredentialBudgetBoundsBlockedConditionalMutation(t *testing.T) {
- gin.SetMode(gin.TestMode)
- account := expiredGrokOAuthAccountForCredentialTest(736)
- baseRepo := &tokenRefreshAccountRepo{}
- baseRepo.accountsByID = map[int64]*Account{account.ID: account}
- repo := &grokCredentialBlockingRepo{
- tokenRefreshAccountRepo: baseRepo,
- setErrorStarted: make(chan struct{}),
- setTempStarted: make(chan struct{}),
- }
- cache := &grokTokenCacheForProviderTest{lockResult: true}
- provider := NewGrokTokenProvider(repo, cache)
- provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{
- err: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_grant"),
- })
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- c.Set(grokCredentialFailoverDeadlineKey, time.Now().Add(40*time.Millisecond))
- startedAt := time.Now()
- token, kind, err := svc.getRequestCredential(context.Background(), c, account)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, GatewayFailureScopeRequest, failoverErr.Scope)
- require.Equal(t, GrokCredentialReasonFailoverTimeout, failoverErr.Reason)
- require.Equal(t, NextAccountStop, failoverErr.NextAccountAction)
- require.Empty(t, token)
- require.Empty(t, kind)
- require.Less(t, time.Since(startedAt), 500*time.Millisecond)
- require.Zero(t, baseRepo.setErrorCalls)
- require.Zero(t, baseRepo.setTempUnschedCalls)
- require.Empty(t, cache.deletedKeys)
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- }
- func TestGetRequestCredentialLockHeldTimeoutDoesNotQuarantineAccount(t *testing.T) {
- gin.SetMode(gin.TestMode)
- tests := []struct {
- name string
- buildRepo func(*Account) AccountRepository
- wantScope GatewayFailureScope
- wantReason GatewayFailureReason
- wantAction NextAccountAction
- }{
- {
- name: "authoritative row unchanged",
- buildRepo: func(account *Account) AccountRepository {
- repo := &tokenRefreshAccountRepo{}
- repo.accountsByID = map[int64]*Account{account.ID: account}
- return repo
- },
- wantScope: GatewayFailureScopeAccount,
- wantReason: GrokCredentialReasonAccountChanged,
- wantAction: NextAccountRetry,
- },
- {
- name: "selected account was deleted",
- buildRepo: func(account *Account) AccountRepository {
- base := &tokenRefreshAccountRepo{}
- base.accountsByID = map[int64]*Account{account.ID: account}
- return &grokCredentialRereadFailureRepo{tokenRefreshAccountRepo: base}
- },
- wantScope: GatewayFailureScopeAccount,
- wantReason: GrokCredentialReasonAccountChanged,
- wantAction: NextAccountRetry,
- },
- {
- name: "shared account store unavailable",
- buildRepo: func(account *Account) AccountRepository {
- base := &tokenRefreshAccountRepo{}
- base.accountsByID = map[int64]*Account{account.ID: account}
- return &grokCredentialRereadFailureRepo{tokenRefreshAccountRepo: base, err: errors.New("database unavailable")}
- },
- wantScope: GatewayFailureScopeProvider,
- wantReason: GrokCredentialReasonProviderDown,
- wantAction: NextAccountStop,
- },
- }
- for index, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(int64(7400 + index))
- repo := tt.buildRepo(account)
- cache := &grokTokenCacheForProviderTest{lockResult: false}
- provider := NewGrokTokenProvider(repo, cache)
- provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{})
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- startedAt := time.Now()
- _, _, err := svc.getRequestCredential(context.Background(), c, account)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, tt.wantScope, failoverErr.Scope)
- require.Equal(t, tt.wantReason, failoverErr.Reason)
- require.Equal(t, tt.wantAction, failoverErr.NextAccountAction)
- require.Less(t, time.Since(startedAt), 3*time.Second)
- switch countingRepo := repo.(type) {
- case *tokenRefreshAccountRepo:
- require.Zero(t, countingRepo.setErrorCalls)
- require.Zero(t, countingRepo.setTempUnschedCalls)
- case *grokCredentialRereadFailureRepo:
- require.Zero(t, countingRepo.tokenRefreshAccountRepo.setErrorCalls)
- require.Zero(t, countingRepo.tokenRefreshAccountRepo.setTempUnschedCalls)
- }
- require.Empty(t, cache.deletedKeys)
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- })
- }
- }
- func TestGrokCredentialMutationCancellationAmbiguityConfirmsDurableCommit(t *testing.T) {
- tests := []struct {
- name string
- class grokCredentialFailureClass
- committed func(*Account) bool
- }{
- {
- name: "permanent quarantine",
- class: grokCredentialFailureClass{reason: GrokCredentialReasonRevoked, permanent: true},
- committed: func(account *Account) bool {
- return account.Status == StatusError && !account.Schedulable && account.ErrorMessage == string(GrokCredentialReasonRevoked)
- },
- },
- {
- name: "temporary quarantine",
- class: grokCredentialFailureClass{reason: GrokCredentialReasonRefreshTransient, transient: true},
- committed: func(account *Account) bool {
- return account.TempUnschedulableUntil != nil && account.TempUnschedulableReason == string(GrokCredentialReasonRefreshTransient)
- },
- },
- }
- for index, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(int64(737 + index))
- baseRepo := &tokenRefreshAccountRepo{}
- baseRepo.accountsByID = map[int64]*Account{account.ID: account}
- repo := &grokCredentialCommitThenCancelRepo{tokenRefreshAccountRepo: baseRepo}
- cache := &grokTokenCacheForProviderTest{}
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, cache)}
- ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
- defer cancel()
- token, err := svc.applyGrokCredentialAccountFailure(ctx, account, tt.class)
- require.ErrorIs(t, err, context.DeadlineExceeded)
- require.Empty(t, token)
- require.True(t, tt.committed(account), "the detached confirmation must recognize the durable mutation")
- require.True(t, svc.isOpenAIAccountRuntimeBlocked(account), "a confirmed durable quarantine must retain its runtime block")
- if tt.class.permanent {
- require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys)
- } else {
- require.Empty(t, cache.deletedKeys)
- }
- })
- }
- }
- func TestGrokCredentialInnerStateDeadlineAmbiguityConfirmsDurableCommit(t *testing.T) {
- tests := []struct {
- name string
- class grokCredentialFailureClass
- }{
- {name: "permanent quarantine", class: grokCredentialFailureClass{reason: GrokCredentialReasonRevoked, permanent: true}},
- {name: "temporary quarantine", class: grokCredentialFailureClass{reason: GrokCredentialReasonRefreshTransient, transient: true}},
- }
- for index, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(int64(7500 + index))
- baseRepo := &tokenRefreshAccountRepo{}
- baseRepo.accountsByID = map[int64]*Account{account.ID: account}
- repo := &grokCredentialCommitThenCancelRepo{
- tokenRefreshAccountRepo: baseRepo,
- returnErr: context.DeadlineExceeded,
- }
- cache := &grokTokenCacheForProviderTest{}
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, cache)}
- token, err := svc.applyGrokCredentialAccountFailure(context.Background(), account, tt.class)
- require.NoError(t, err, "the detached readback must resolve the inner timeout's commit ambiguity")
- require.Empty(t, token)
- require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
- if tt.class.permanent {
- require.Equal(t, StatusError, account.Status)
- require.False(t, account.Schedulable)
- require.Equal(t, []string{GrokTokenCacheKey(account)}, cache.deletedKeys)
- } else {
- require.NotNil(t, account.TempUnschedulableUntil)
- require.Equal(t, string(GrokCredentialReasonRefreshTransient), account.TempUnschedulableReason)
- require.Empty(t, cache.deletedKeys)
- }
- })
- }
- }
- func TestGrokCredentialUnconfirmedInnerStateDeadlineStopsAndRetainsSafetyBlock(t *testing.T) {
- tests := []struct {
- name string
- class grokCredentialFailureClass
- }{
- {name: "permanent quarantine", class: grokCredentialFailureClass{reason: GrokCredentialReasonRevoked, permanent: true}},
- {name: "temporary quarantine", class: grokCredentialFailureClass{reason: GrokCredentialReasonRefreshTransient, transient: true}},
- }
- for index, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(int64(7600 + index))
- baseRepo := &tokenRefreshAccountRepo{}
- baseRepo.accountsByID = map[int64]*Account{account.ID: account}
- repo := &grokCredentialUncommittedDeadlineRepo{tokenRefreshAccountRepo: baseRepo}
- cache := &grokTokenCacheForProviderTest{}
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, cache)}
- token, err := svc.applyGrokCredentialAccountFailure(context.Background(), account, tt.class)
- require.ErrorIs(t, err, errGrokCredentialStateUpdateFailed)
- require.Empty(t, token)
- require.True(t, svc.isOpenAIAccountRuntimeBlocked(account), "an unknown commit outcome must retain the local safety block")
- require.Equal(t, StatusActive, account.Status)
- require.True(t, account.Schedulable)
- require.Nil(t, account.TempUnschedulableUntil)
- require.Empty(t, cache.deletedKeys)
- })
- }
- }
- func TestGrokCredentialRuntimeRollbackOwnership(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(734)
- t.Run("later extending block survives", func(t *testing.T) {
- svc := &OpenAIGatewayService{}
- until := time.Now().Add(time.Minute)
- rollbackFirst := svc.blockGrokCredentialRuntime(account, until, "first")
- secondInstalled := make(chan struct{})
- go func() {
- svc.BlockAccountScheduling(account, until.Add(time.Minute), "independent")
- close(secondInstalled)
- }()
- <-secondInstalled
- rollbackFirst()
- require.True(t, svc.isOpenAIAccountRuntimeBlocked(account),
- "rollback owned by the first invocation must not remove a later extending block")
- })
- t.Run("independent shorter block steals rollback ownership", func(t *testing.T) {
- svc := &OpenAIGatewayService{}
- until := time.Now().Add(2 * time.Minute)
- rollbackFirst := svc.blockGrokCredentialRuntime(account, until, "first")
- svc.BlockAccountScheduling(account, until.Add(-time.Minute), "shorter-no-op")
- rollbackFirst()
- require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
- })
- t.Run("serialized tentative rollbacks leave no block", func(t *testing.T) {
- svc := &OpenAIGatewayService{}
- for i := 0; i < 2; i++ {
- mu := svc.grokCredentialMutationLock(account.ID)
- require.NoError(t, mu.Lock(context.Background()))
- rollback := svc.blockGrokCredentialRuntime(account, time.Now().Add(time.Minute), "tentative")
- rollback()
- mu.Unlock()
- }
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- })
- }
- func TestGetRequestCredentialAPIKeyBypassesOAuthFailureMapping(t *testing.T) {
- gin.SetMode(gin.TestMode)
- account := &Account{
- ID: 705,
- Platform: PlatformGrok,
- Type: AccountTypeAPIKey,
- Credentials: map[string]any{
- "api_key": "third-party-key",
- "base_url": "https://grok.example.test/v1",
- },
- }
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- svc := &OpenAIGatewayService{}
- token, kind, err := svc.getRequestCredential(context.Background(), c, account)
- require.NoError(t, err)
- require.Equal(t, "third-party-key", token)
- require.Equal(t, "apikey", kind)
- _, hasEvents := c.Get(OpsUpstreamErrorsKey)
- require.False(t, hasEvents)
- }
- func TestPermanentCredentialFailureDoesNotDisableConcurrentlyRefreshedAccount(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(707)
- latest := *account
- latest.Credentials = shallowCopyMap(account.Credentials)
- latest.Credentials["access_token"] = "fresh-access-token"
- latest.Credentials["refresh_token"] = "rotated-refresh-token"
- latest.Credentials["expires_at"] = time.Now().Add(time.Hour).UTC().Format(time.RFC3339)
- latest.Credentials["_token_version"] = time.Now().UnixMilli()
- repo := &tokenRefreshAccountRepo{}
- repo.accountsByID = map[int64]*Account{account.ID: &latest}
- cache := &grokTokenCacheForProviderTest{}
- svc := &OpenAIGatewayService{
- accountRepo: repo,
- grokTokenProvider: NewGrokTokenProvider(repo, cache),
- }
- _, mutationErr := svc.applyGrokCredentialAccountFailure(context.Background(), account, grokCredentialFailureClass{
- scope: GatewayFailureScopeAccount,
- reason: GrokCredentialReasonRevoked,
- action: NextAccountRetry,
- permanent: true,
- })
- require.NoError(t, mutationErr)
- require.Zero(t, repo.setErrorCalls)
- require.Empty(t, cache.deletedKeys)
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- }
- func TestCredentialFailureConditionalMutationLosesToConcurrentRefresh(t *testing.T) {
- for index, tt := range []struct {
- name string
- refreshErr error
- }{
- {name: "permanent", refreshErr: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_grant")},
- {name: "transient", refreshErr: errors.New("temporary refresh transport failure")},
- } {
- t.Run(tt.name, func(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(int64(770 + index))
- repo := &tokenRefreshAccountRepo{}
- repo.accountsByID = map[int64]*Account{account.ID: account}
- repo.beforeConditionalState = func() {
- fresh := *account
- fresh.Credentials = shallowCopyMap(account.Credentials)
- fresh.Credentials["access_token"] = "refresh-won-token"
- fresh.Credentials["refresh_token"] = "refresh-won-refresh"
- fresh.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339)
- fresh.Credentials["_token_version"] = time.Now().UnixMilli()
- repo.accountsByID[account.ID] = &fresh
- }
- cache := &grokTokenCacheForProviderTest{lockResult: true}
- provider := NewGrokTokenProvider(repo, cache)
- provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{err: tt.refreshErr})
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- token, kind, err := svc.getRequestCredential(context.Background(), c, account)
- require.NoError(t, err)
- require.Equal(t, "refresh-won-token", token)
- require.Equal(t, "oauth", kind)
- require.Zero(t, repo.setErrorCalls)
- require.Zero(t, repo.setTempUnschedCalls)
- require.Empty(t, cache.deletedKeys)
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- })
- }
- }
- func TestCredentialFailureConditionalMutationLosesToConcurrentProxyRepair(t *testing.T) {
- for index, tt := range []struct {
- name string
- refreshErr error
- }{
- {name: "permanent", refreshErr: infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_grant")},
- {name: "transient", refreshErr: errors.New("temporary refresh transport failure")},
- } {
- t.Run(tt.name, func(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(int64(780 + index))
- oldProxyID := int64(10)
- account.ProxyID = &oldProxyID
- account.Proxy = &Proxy{}
- repo := &tokenRefreshAccountRepo{}
- repo.accountsByID = map[int64]*Account{account.ID: account}
- repo.beforeConditionalState = func() {
- fresh := *account
- fresh.Credentials = shallowCopyMap(account.Credentials)
- repairedProxyID := int64(11)
- fresh.ProxyID = &repairedProxyID
- repo.accountsByID[account.ID] = &fresh
- }
- cache := &grokTokenCacheForProviderTest{lockResult: true}
- provider := NewGrokTokenProvider(repo, cache)
- provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{err: tt.refreshErr})
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- _, _, err := svc.getRequestCredential(context.Background(), c, account)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, GrokCredentialReasonAccountChanged, failoverErr.Reason)
- require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
- require.Zero(t, repo.setErrorCalls)
- require.Zero(t, repo.setTempUnschedCalls)
- require.Empty(t, cache.deletedKeys)
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- })
- }
- }
- func TestCredentialFailureConditionalMutationLosesToSameIDProxyRestoration(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(790)
- proxyID := int64(10)
- account.ProxyID = &proxyID
- account.Proxy = nil
- repo := &tokenRefreshAccountRepo{}
- repo.accountsByID = map[int64]*Account{account.ID: account}
- repo.beforeConditionalState = func() {
- account.Proxy = &Proxy{ID: proxyID}
- }
- cache := &grokTokenCacheForProviderTest{lockResult: true}
- provider := NewGrokTokenProvider(repo, cache)
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- _, _, err := svc.getRequestCredential(context.Background(), c, account)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, GrokCredentialReasonAccountChanged, failoverErr.Reason)
- require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
- require.Zero(t, repo.setErrorCalls)
- require.Zero(t, repo.setTempUnschedCalls)
- require.Empty(t, cache.deletedKeys)
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- }
- func TestCredentialFailureConditionalMutationLosesToConcurrentUnschedulableState(t *testing.T) {
- future := time.Now().Add(time.Hour)
- states := []struct {
- name string
- mutate func(*Account)
- }{
- {name: "admin schedulable false", mutate: func(account *Account) { account.Schedulable = false }},
- {name: "temporary cooldown", mutate: func(account *Account) { account.TempUnschedulableUntil = &future }},
- {name: "rate limit cooldown", mutate: func(account *Account) { account.RateLimitResetAt = &future }},
- {name: "overload cooldown", mutate: func(account *Account) { account.OverloadUntil = &future }},
- }
- classes := []struct {
- name string
- class grokCredentialFailureClass
- }{
- {name: "permanent", class: grokCredentialFailureClass{reason: GrokCredentialReasonRevoked, permanent: true}},
- {name: "transient", class: grokCredentialFailureClass{reason: GrokCredentialReasonRefreshTransient, transient: true}},
- }
- for classIndex, classCase := range classes {
- for stateIndex, stateCase := range states {
- t.Run(classCase.name+"/"+stateCase.name, func(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(int64(791 + classIndex*10 + stateIndex))
- repo := &tokenRefreshAccountRepo{}
- repo.accountsByID = map[int64]*Account{account.ID: account}
- repo.beforeConditionalState = func() { stateCase.mutate(account) }
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, &grokTokenCacheForProviderTest{})}
- token, err := svc.applyGrokCredentialAccountFailure(context.Background(), account, classCase.class)
- require.ErrorIs(t, err, errOAuthRefreshAccountStateChanged)
- require.Empty(t, token)
- require.Zero(t, repo.setErrorCalls)
- require.Zero(t, repo.setTempUnschedCalls)
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- })
- }
- }
- }
- func TestCredentialFailureCASMissDoesNotRecoverIneligibleLatestCredential(t *testing.T) {
- future := time.Now().Add(time.Hour)
- states := []struct {
- name string
- mutate func(*OpenAIGatewayService, *Account)
- wantRuntimeBlocked bool
- }{
- {name: "disabled", mutate: func(_ *OpenAIGatewayService, account *Account) { account.Status = StatusDisabled }},
- {name: "not schedulable", mutate: func(_ *OpenAIGatewayService, account *Account) { account.Schedulable = false }},
- {name: "temporarily unschedulable", mutate: func(_ *OpenAIGatewayService, account *Account) { account.TempUnschedulableUntil = &future }},
- {name: "rate limited", mutate: func(_ *OpenAIGatewayService, account *Account) { account.RateLimitResetAt = &future }},
- {name: "overloaded", mutate: func(_ *OpenAIGatewayService, account *Account) { account.OverloadUntil = &future }},
- {
- name: "independently runtime blocked",
- mutate: func(svc *OpenAIGatewayService, account *Account) {
- svc.BlockAccountScheduling(account, time.Now().Add(24*time.Hour), "independent")
- },
- wantRuntimeBlocked: true,
- },
- }
- classes := []struct {
- name string
- class grokCredentialFailureClass
- }{
- {name: "permanent", class: grokCredentialFailureClass{reason: GrokCredentialReasonRevoked, permanent: true}},
- {name: "transient", class: grokCredentialFailureClass{reason: GrokCredentialReasonRefreshTransient, transient: true}},
- }
- for classIndex, classCase := range classes {
- for stateIndex, stateCase := range states {
- t.Run(classCase.name+"/"+stateCase.name, func(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(int64(800 + classIndex*20 + stateIndex))
- repo := &tokenRefreshAccountRepo{}
- repo.accountsByID = map[int64]*Account{account.ID: account}
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, &grokTokenCacheForProviderTest{})}
- repo.beforeConditionalState = func() {
- latest := *account
- latest.Credentials = shallowCopyMap(account.Credentials)
- latest.Credentials["access_token"] = "fresh-but-ineligible-token"
- latest.Credentials["refresh_token"] = "fresh-but-ineligible-refresh"
- latest.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339)
- latest.Credentials["_token_version"] = time.Now().UnixMilli()
- stateCase.mutate(svc, &latest)
- repo.accountsByID[account.ID] = &latest
- }
- token, err := svc.applyGrokCredentialAccountFailure(context.Background(), account, classCase.class)
- require.ErrorIs(t, err, errOAuthRefreshAccountStateChanged)
- require.Empty(t, token)
- require.Zero(t, repo.setErrorCalls)
- require.Zero(t, repo.setTempUnschedCalls)
- require.Equal(t, stateCase.wantRuntimeBlocked, svc.isOpenAIAccountRuntimeBlocked(account))
- })
- }
- }
- }
- func TestGetRequestCredentialSharedCredentialPersistenceFailureStopsWithoutAccountMutation(t *testing.T) {
- account := expiredGrokOAuthAccountForCredentialTest(782)
- repo := &tokenRefreshAccountRepo{conditionalSuccessErr: errors.New("database unavailable")}
- repo.accountsByID = map[int64]*Account{account.ID: account}
- cache := &grokTokenCacheForProviderTest{lockResult: true}
- provider := NewGrokTokenProvider(repo, cache)
- provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{credentials: map[string]any{
- "access_token": "new-access-token",
- "refresh_token": "new-refresh-token",
- "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
- }})
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- _, _, err := svc.getRequestCredential(context.Background(), c, account)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, GatewayFailureScopeProvider, failoverErr.Scope)
- require.Equal(t, GrokCredentialReasonProviderDown, failoverErr.Reason)
- require.Equal(t, NextAccountStop, failoverErr.NextAccountAction)
- require.Equal(t, 1, repo.conditionalSuccessCalls)
- require.Zero(t, repo.updateCredentialsCalls)
- require.Zero(t, repo.setErrorCalls)
- require.Zero(t, repo.setTempUnschedCalls)
- require.Empty(t, cache.deletedKeys)
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- }
- func TestGetRequestCredentialRecoversConcurrentRefreshWithoutFailover(t *testing.T) {
- gin.SetMode(gin.TestMode)
- account := expiredGrokOAuthAccountForCredentialTest(733)
- latest := *account
- latest.Credentials = shallowCopyMap(account.Credentials)
- latest.Credentials["access_token"] = "fresh-concurrent-access"
- latest.Credentials["refresh_token"] = "fresh-concurrent-refresh"
- latest.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339)
- latest.Credentials["_token_version"] = time.Now().UnixMilli()
- baseRepo := &tokenRefreshAccountRepo{}
- baseRepo.accountsByID = map[int64]*Account{account.ID: account}
- repo := &grokCredentialSequencedRepo{tokenRefreshAccountRepo: baseRepo, latest: &latest}
- cache := &grokTokenCacheForProviderTest{lockResult: true}
- provider := NewGrokTokenProvider(repo, cache)
- provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{
- err: infraerrors.New(http.StatusForbidden, "GROK_OAUTH_ENTITLEMENT_DENIED", "access_denied"),
- })
- svc := &OpenAIGatewayService{accountRepo: repo, grokTokenProvider: provider}
- c, _ := gin.CreateTestContext(httptest.NewRecorder())
- token, kind, err := svc.getRequestCredential(context.Background(), c, account)
- require.NoError(t, err)
- require.Equal(t, "fresh-concurrent-access", token)
- require.Equal(t, "oauth", kind)
- require.Zero(t, baseRepo.setErrorCalls)
- require.Zero(t, baseRepo.setTempUnschedCalls)
- require.Empty(t, cache.deletedKeys)
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- _, hasEvents := c.Get(OpsUpstreamErrorsKey)
- require.False(t, hasEvents)
- }
- func expiredGrokOAuthAccountForCredentialTest(id int64) *Account {
- return &Account{
- ID: id,
- Platform: PlatformGrok,
- Type: AccountTypeOAuth,
- Status: StatusActive,
- Schedulable: true,
- Credentials: map[string]any{
- "access_token": "expired-access-token",
- "refresh_token": "refresh-token",
- "expires_at": time.Now().Add(-time.Minute).UTC().Format(time.RFC3339),
- "base_url": xai.DefaultCLIBaseURL,
- },
- }
- }
|