grok_credential_failure.go 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670
  1. package service
  2. import (
  3. "context"
  4. "encoding/json"
  5. "errors"
  6. "fmt"
  7. "log/slog"
  8. "net/http"
  9. "strings"
  10. "time"
  11. infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
  12. "github.com/gin-gonic/gin"
  13. )
  14. const (
  15. grokCredentialFailoverDeadlineKey = "grok_credential_failover_deadline"
  16. grokCredentialFailoverBudget = 15 * time.Second
  17. grokCredentialMutationTimeout = 5 * time.Second
  18. grokCredentialMutationConfirmWait = 250 * time.Millisecond
  19. grokCredentialCacheCleanupTimeout = 500 * time.Millisecond
  20. GrokCredentialUnavailableClientMessage = "No healthy Grok OAuth account is currently available"
  21. GrokCredentialReasonRevoked GatewayFailureReason = "grok_oauth_credential_revoked"
  22. GrokCredentialReasonMissing GatewayFailureReason = "grok_oauth_credentials_missing"
  23. GrokCredentialReasonEntitlement GatewayFailureReason = "grok_oauth_entitlement_action_required"
  24. GrokCredentialReasonProxyInvalid GatewayFailureReason = "grok_oauth_proxy_invalid"
  25. GrokCredentialReasonRefreshTransient GatewayFailureReason = "grok_oauth_refresh_transient"
  26. GrokCredentialReasonProviderConfig GatewayFailureReason = "grok_oauth_provider_config"
  27. GrokCredentialReasonProviderDown GatewayFailureReason = "grok_oauth_provider_unavailable"
  28. GrokCredentialReasonAccountChanged GatewayFailureReason = "grok_oauth_account_state_changed"
  29. GrokCredentialReasonStateUpdate GatewayFailureReason = "grok_oauth_account_state_update_failed"
  30. GrokCredentialReasonFailoverTimeout GatewayFailureReason = "grok_oauth_failover_timeout"
  31. )
  32. var errGrokCredentialStateUpdateFailed = errors.New("grok oauth account state update failed")
  33. type grokCredentialFailureClass struct {
  34. scope GatewayFailureScope
  35. reason GatewayFailureReason
  36. action NextAccountAction
  37. permanent bool
  38. transient bool
  39. message string
  40. snapshot *GrokCredentialMutationSnapshot
  41. }
  42. // GrokCredentialMutationSnapshot is the credential identity observed when the
  43. // request selected an account. Repository mutations compare all fields before
  44. // quarantining that account so a concurrent refresh cannot be overwritten.
  45. type GrokCredentialMutationSnapshot struct {
  46. CredentialsJSON string
  47. AccessToken string
  48. RefreshToken string
  49. TokenVersion int64
  50. ProxyID *int64
  51. }
  52. type grokCredentialFailureSnapshotError struct {
  53. cause error
  54. snapshot GrokCredentialMutationSnapshot
  55. }
  56. func (e *grokCredentialFailureSnapshotError) Error() string { return e.cause.Error() }
  57. func (e *grokCredentialFailureSnapshotError) Unwrap() error { return e.cause }
  58. func withGrokCredentialFailureSnapshot(err error, account *Account) error {
  59. if err == nil || account == nil || !account.IsGrokOAuth() {
  60. return err
  61. }
  62. var existing *grokCredentialFailureSnapshotError
  63. if errors.As(err, &existing) {
  64. return err
  65. }
  66. return &grokCredentialFailureSnapshotError{cause: err, snapshot: grokCredentialMutationSnapshot(account)}
  67. }
  68. func grokCredentialFailureSnapshot(err error) (GrokCredentialMutationSnapshot, bool) {
  69. var snapshotErr *grokCredentialFailureSnapshotError
  70. if !errors.As(err, &snapshotErr) || snapshotErr == nil {
  71. return GrokCredentialMutationSnapshot{}, false
  72. }
  73. return snapshotErr.snapshot, true
  74. }
  75. type grokCredentialConditionalStateRepository interface {
  76. SetGrokCredentialErrorIfMatch(context.Context, int64, GrokCredentialMutationSnapshot, string) (bool, error)
  77. SetGrokCredentialTempUnschedulableIfMatch(context.Context, int64, GrokCredentialMutationSnapshot, time.Time, string) (bool, error)
  78. }
  79. // GetRequestCredential applies the request-path credential and failover contract
  80. // before any upstream transport is opened.
  81. func (s *OpenAIGatewayService) GetRequestCredential(ctx context.Context, c *gin.Context, account *Account) (string, string, error) {
  82. return s.getRequestCredential(ctx, c, account)
  83. }
  84. func (s *OpenAIGatewayService) getRequestCredential(ctx context.Context, c *gin.Context, account *Account) (string, string, error) {
  85. if ctx == nil {
  86. ctx = context.Background()
  87. }
  88. if account == nil {
  89. return "", "", errors.New("account is nil")
  90. }
  91. if !account.IsGrokOAuth() {
  92. return s.GetAccessToken(ctx, account)
  93. }
  94. if err := ctx.Err(); err != nil {
  95. return "", "", err
  96. }
  97. if s == nil || s.grokTokenProvider == nil {
  98. return "", "", s.newGrokCredentialFailover(c, account, grokCredentialFailureClass{
  99. scope: GatewayFailureScopeProvider,
  100. reason: GrokCredentialReasonProviderConfig,
  101. action: NextAccountStop,
  102. message: "Grok OAuth credential provider is unavailable",
  103. })
  104. }
  105. if s.isOpenAIAccountRuntimeBlocked(account) {
  106. return "", "", s.newGrokCredentialFailover(c, account, grokCredentialFailureClass{
  107. scope: GatewayFailureScopeAccount,
  108. reason: GrokCredentialReasonAccountChanged,
  109. action: NextAccountRetry,
  110. message: "Grok OAuth account is not currently schedulable",
  111. })
  112. }
  113. credentialCtx, cancel, budgetExpired := grokCredentialAcquisitionContext(ctx, c)
  114. if cancel != nil {
  115. defer cancel()
  116. }
  117. if budgetExpired {
  118. return "", "", s.newGrokCredentialFailover(c, account, grokCredentialFailureClass{
  119. scope: GatewayFailureScopeRequest,
  120. reason: GrokCredentialReasonFailoverTimeout,
  121. action: NextAccountStop,
  122. message: "Grok OAuth credential failover budget exhausted",
  123. })
  124. }
  125. token, kind, err := s.GetAccessToken(credentialCtx, account)
  126. if err == nil {
  127. if s.isOpenAIAccountRuntimeBlocked(account) {
  128. return "", "", s.newGrokCredentialFailover(c, account, grokCredentialFailureClass{
  129. scope: GatewayFailureScopeAccount,
  130. reason: GrokCredentialReasonAccountChanged,
  131. action: NextAccountRetry,
  132. message: "Grok OAuth account is not currently schedulable",
  133. })
  134. }
  135. return token, kind, nil
  136. }
  137. if parentErr := ctx.Err(); parentErr != nil {
  138. return "", "", parentErr
  139. }
  140. if credentialCtx.Err() != nil {
  141. return "", "", s.newGrokCredentialFailover(c, account, grokCredentialFailureClass{
  142. scope: GatewayFailureScopeRequest,
  143. reason: GrokCredentialReasonFailoverTimeout,
  144. action: NextAccountStop,
  145. message: "Grok OAuth credential failover budget exhausted",
  146. })
  147. }
  148. class := classifyGrokCredentialFailure(account, err)
  149. if snapshot, ok := grokCredentialFailureSnapshot(err); ok {
  150. class.snapshot = &snapshot
  151. }
  152. if ctx.Err() != nil {
  153. return "", "", ctx.Err()
  154. }
  155. if class.permanent || class.transient {
  156. freshToken, mutationErr := s.applyGrokCredentialAccountFailure(credentialCtx, account, class)
  157. if freshToken != "" {
  158. return freshToken, "oauth", nil
  159. }
  160. if mutationErr != nil {
  161. if ctx.Err() != nil {
  162. return "", "", ctx.Err()
  163. }
  164. if credentialCtx.Err() != nil {
  165. return "", "", s.newGrokCredentialFailover(c, account, grokCredentialFailureClass{
  166. scope: GatewayFailureScopeRequest,
  167. reason: GrokCredentialReasonFailoverTimeout,
  168. action: NextAccountStop,
  169. message: "Grok OAuth credential failover budget exhausted",
  170. })
  171. }
  172. if errors.Is(mutationErr, errOAuthRefreshAccountStateChanged) {
  173. class = grokCredentialFailureClass{
  174. scope: GatewayFailureScopeAccount,
  175. reason: GrokCredentialReasonAccountChanged,
  176. action: NextAccountRetry,
  177. message: "Grok OAuth account eligibility changed",
  178. }
  179. } else if errors.Is(mutationErr, errOAuthRefreshAccountRereadFailed) {
  180. class = grokCredentialFailureClass{
  181. scope: GatewayFailureScopeProvider,
  182. reason: GrokCredentialReasonProviderDown,
  183. action: NextAccountStop,
  184. message: "Grok OAuth account state is temporarily unavailable",
  185. }
  186. } else {
  187. class = grokCredentialFailureClass{
  188. scope: GatewayFailureScopeProvider,
  189. reason: GrokCredentialReasonStateUpdate,
  190. action: NextAccountStop,
  191. message: "Grok OAuth account state could not be updated safely",
  192. }
  193. }
  194. }
  195. }
  196. return "", "", s.newGrokCredentialFailover(c, account, class)
  197. }
  198. func grokCredentialAcquisitionContext(ctx context.Context, c *gin.Context) (context.Context, context.CancelFunc, bool) {
  199. if c == nil {
  200. return ctx, nil, false
  201. }
  202. deadline := time.Time{}
  203. if raw, ok := c.Get(grokCredentialFailoverDeadlineKey); ok {
  204. deadline, _ = raw.(time.Time)
  205. }
  206. if deadline.IsZero() {
  207. deadline = time.Now().Add(grokCredentialFailoverBudget)
  208. c.Set(grokCredentialFailoverDeadlineKey, deadline)
  209. }
  210. if !time.Now().Before(deadline) {
  211. return ctx, nil, true
  212. }
  213. acquireCtx, cancel := context.WithDeadline(ctx, deadline)
  214. return acquireCtx, cancel, false
  215. }
  216. func classifyGrokCredentialFailure(account *Account, err error) grokCredentialFailureClass {
  217. stableReason := strings.ToLower(strings.TrimSpace(infraerrors.Reason(err)))
  218. message := ""
  219. if err != nil {
  220. message = strings.ToLower(err.Error())
  221. }
  222. contains := func(values ...string) bool {
  223. for _, value := range values {
  224. if strings.Contains(stableReason, value) || strings.Contains(message, value) {
  225. return true
  226. }
  227. }
  228. return false
  229. }
  230. var providerConfigErr *providerConfigurationRefreshError
  231. var containmentErr *providerCycleContainmentRefreshError
  232. switch {
  233. case errors.Is(err, errGrokOAuthRefreshTokenMissing), errors.Is(err, errGrokOAuthAccessTokenMissing), errors.Is(err, errGrokOAuthAccessTokenExpired):
  234. return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonMissing, action: NextAccountRetry, permanent: true, message: "Grok OAuth credentials are missing or expired"}
  235. case contains("invalid_grant", "invalid_refresh_token", "token_expired", "refresh_token_reused", "refresh_token_invalidated", "app_session_terminated"):
  236. return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonRevoked, action: NextAccountRetry, permanent: true, message: "Grok OAuth credentials require account action"}
  237. case contains("spending limit", "run out of credits", "out of credits", "credits exhausted", "included free usage"):
  238. // Billing and rolling free-usage exhaustion recover without replacing the
  239. // OAuth credential. Treat refresh failures as transient so the account
  240. // remains eligible for a later quota probe.
  241. return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonRefreshTransient, action: NextAccountRetry, transient: true, message: "Grok OAuth billing quota is temporarily exhausted"}
  242. case contains("grok_oauth_entitlement_denied", "entitlement_denied", "access_denied", "subscription required", "no active grok subscription"):
  243. return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonEntitlement, action: NextAccountRetry, permanent: true, message: "Grok OAuth entitlement requires account action"}
  244. case errors.Is(err, errGrokOAuthConfiguredProxyMiss), contains("grok_oauth_proxy_not_found"):
  245. return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonProxyInvalid, action: NextAccountRetry, permanent: true, message: "Grok OAuth account proxy configuration is invalid"}
  246. case errors.Is(err, errOAuthRefreshAccountRereadFailed):
  247. return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderDown, action: NextAccountStop, message: "Grok OAuth account state is temporarily unavailable"}
  248. case errors.Is(err, errOAuthRefreshCredentialPersist):
  249. return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderDown, action: NextAccountStop, message: "Grok OAuth shared credential state is temporarily unavailable"}
  250. case errors.As(err, &containmentErr):
  251. return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderDown, action: NextAccountStop, message: "Grok OAuth provider state is temporarily unavailable"}
  252. case errors.Is(err, errOAuthRefreshAccountStateChanged):
  253. return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonAccountChanged, action: NextAccountRetry, message: "Grok OAuth account eligibility changed"}
  254. case errors.As(err, &providerConfigErr):
  255. return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderConfig, action: NextAccountStop, message: "Grok OAuth provider configuration is unavailable"}
  256. case errors.Is(err, errGrokOAuthRefreshNotConfigured), contains("invalid_client", "unauthorized_client", "invalid_scope", "unknown scope", "grok oauth service is not configured", "grok_oauth_proxy_not_available"):
  257. return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderConfig, action: NextAccountStop, message: "Grok OAuth provider configuration is unavailable"}
  258. case contains("grok_oauth_proxy_lookup_failed"),
  259. contains("grok_oauth_token_refresh_failed") && contains("status 403") && (account == nil || account.ProxyID == nil):
  260. return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderDown, action: NextAccountStop, message: "Grok OAuth provider is temporarily unavailable"}
  261. case contains("grok_oauth_client_init_failed") && (account == nil || account.ProxyID == nil):
  262. return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderConfig, action: NextAccountStop, message: "Grok OAuth provider configuration is unavailable"}
  263. case contains("grok_oauth_request_failed") && (account == nil || account.ProxyID == nil):
  264. return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderDown, action: NextAccountStop, message: "Grok OAuth provider is temporarily unavailable"}
  265. case contains("status 429", "status 500", "status 502", "status 503", "status 504") && (account == nil || account.ProxyID == nil):
  266. return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderDown, action: NextAccountStop, message: "Grok OAuth provider is temporarily unavailable"}
  267. default:
  268. return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonRefreshTransient, action: NextAccountRetry, transient: true, message: "Grok OAuth credential refresh is temporarily unavailable"}
  269. }
  270. }
  271. func (s *OpenAIGatewayService) applyGrokCredentialAccountFailure(ctx context.Context, account *Account, class grokCredentialFailureClass) (string, error) {
  272. if s == nil || account == nil || ctx == nil || ctx.Err() != nil {
  273. if ctx != nil {
  274. return "", ctx.Err()
  275. }
  276. return "", context.Canceled
  277. }
  278. mutationMu := s.grokCredentialMutationLock(account.ID)
  279. if err := mutationMu.Lock(ctx); err != nil {
  280. return "", err
  281. }
  282. defer mutationMu.Unlock()
  283. stateRepo, hasConditionalStateRepo := s.accountRepo.(grokCredentialConditionalStateRepository)
  284. snapshot := grokCredentialMutationSnapshot(account)
  285. if class.snapshot != nil {
  286. snapshot = *class.snapshot
  287. }
  288. if token, err := s.validateCurrentGrokCredentialFailure(ctx, account.ID, snapshot, class); err != nil || token != "" {
  289. return token, err
  290. }
  291. if class.permanent {
  292. if token, ok := s.grokCredentialConcurrentlyRefreshedToken(ctx, account.ID, snapshot); ok {
  293. return token, nil
  294. }
  295. if ctx.Err() != nil {
  296. return "", ctx.Err()
  297. }
  298. rollbackRuntime := s.blockGrokCredentialRuntime(account, time.Time{}, string(class.reason))
  299. keepRuntimeBlock := false
  300. runtimeRollbackDone := false
  301. defer func() {
  302. if !keepRuntimeBlock && !runtimeRollbackDone {
  303. rollbackRuntime()
  304. }
  305. }()
  306. if s.accountRepo == nil {
  307. keepRuntimeBlock = true
  308. return "", fmt.Errorf("%w: account repository is not configured", errGrokCredentialStateUpdateFailed)
  309. }
  310. if !hasConditionalStateRepo {
  311. keepRuntimeBlock = true
  312. return "", fmt.Errorf("%w: conditional account repository is not configured", errGrokCredentialStateUpdateFailed)
  313. }
  314. stateCtx, cancel := context.WithTimeout(ctx, grokCredentialMutationTimeout)
  315. if err := ctx.Err(); err != nil {
  316. cancel()
  317. return "", err
  318. }
  319. updated, err := stateRepo.SetGrokCredentialErrorIfMatch(stateCtx, account.ID, snapshot, string(class.reason))
  320. requestErr := ctx.Err()
  321. cancel()
  322. if err != nil {
  323. slog.Warn("grok_credential_failure.set_error_failed", "account_id", account.ID, "reason", class.reason, "error", err)
  324. if requestErr != nil || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
  325. if s.grokCredentialMutationCommitted(account.ID, class, time.Time{}) {
  326. updated = true
  327. } else if requestErr != nil {
  328. return "", requestErr
  329. } else {
  330. keepRuntimeBlock = true
  331. return "", fmt.Errorf("%w: permanent state commit could not be confirmed: %v", errGrokCredentialStateUpdateFailed, err)
  332. }
  333. } else {
  334. keepRuntimeBlock = true
  335. return "", fmt.Errorf("%w: persist permanent state: %v", errGrokCredentialStateUpdateFailed, err)
  336. }
  337. }
  338. if !updated {
  339. rollbackRuntime()
  340. runtimeRollbackDone = true
  341. return s.resolveGrokCredentialCASMiss(ctx, account.ID, snapshot)
  342. }
  343. // SetError is the linearization point: the durable quarantine is now
  344. // committed and this node's runtime block must not be rolled back.
  345. keepRuntimeBlock = true
  346. if s.grokTokenProvider == nil {
  347. if ctx.Err() != nil {
  348. return "", ctx.Err()
  349. }
  350. return "", fmt.Errorf("%w: token provider is not configured", errGrokCredentialStateUpdateFailed)
  351. }
  352. invalidateCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), grokCredentialCacheCleanupTimeout)
  353. err = s.grokTokenProvider.InvalidateToken(invalidateCtx, account)
  354. cancel()
  355. if err != nil {
  356. slog.Warn("grok_credential_failure.invalidate_token_failed", "account_id", account.ID, "reason", class.reason, "error", err)
  357. if ctx.Err() != nil {
  358. return "", ctx.Err()
  359. }
  360. return "", fmt.Errorf("%w: invalidate cached credential: %v", errGrokCredentialStateUpdateFailed, err)
  361. }
  362. if ctx.Err() != nil {
  363. return "", ctx.Err()
  364. }
  365. return "", nil
  366. }
  367. if class.transient {
  368. until := time.Now().Add(tokenRefreshTempUnschedDuration)
  369. if ctx.Err() != nil {
  370. return "", ctx.Err()
  371. }
  372. rollbackRuntime := s.blockGrokCredentialRuntime(account, until, string(class.reason))
  373. keepRuntimeBlock := false
  374. runtimeRollbackDone := false
  375. defer func() {
  376. if !keepRuntimeBlock && !runtimeRollbackDone {
  377. rollbackRuntime()
  378. }
  379. }()
  380. stateCtx, cancel := context.WithTimeout(ctx, grokCredentialMutationTimeout)
  381. if s.accountRepo == nil {
  382. cancel()
  383. keepRuntimeBlock = true
  384. return "", fmt.Errorf("%w: account repository is not configured", errGrokCredentialStateUpdateFailed)
  385. }
  386. if !hasConditionalStateRepo {
  387. cancel()
  388. keepRuntimeBlock = true
  389. return "", fmt.Errorf("%w: conditional account repository is not configured", errGrokCredentialStateUpdateFailed)
  390. }
  391. if err := ctx.Err(); err != nil {
  392. cancel()
  393. return "", err
  394. }
  395. updated, err := stateRepo.SetGrokCredentialTempUnschedulableIfMatch(stateCtx, account.ID, snapshot, until, string(class.reason))
  396. requestErr := ctx.Err()
  397. cancel()
  398. if err != nil {
  399. slog.Warn("grok_credential_failure.set_temp_unschedulable_failed", "account_id", account.ID, "reason", class.reason, "error", err)
  400. if requestErr != nil || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
  401. if s.grokCredentialMutationCommitted(account.ID, class, until) {
  402. updated = true
  403. } else if requestErr != nil {
  404. return "", requestErr
  405. } else {
  406. keepRuntimeBlock = true
  407. return "", fmt.Errorf("%w: transient state commit could not be confirmed: %v", errGrokCredentialStateUpdateFailed, err)
  408. }
  409. } else {
  410. keepRuntimeBlock = true
  411. return "", fmt.Errorf("%w: persist transient state: %v", errGrokCredentialStateUpdateFailed, err)
  412. }
  413. }
  414. if !updated {
  415. rollbackRuntime()
  416. runtimeRollbackDone = true
  417. return s.resolveGrokCredentialCASMiss(ctx, account.ID, snapshot)
  418. }
  419. // The temporary quarantine is durable after SetTempUnschedulable succeeds.
  420. keepRuntimeBlock = true
  421. if ctx.Err() != nil {
  422. return "", ctx.Err()
  423. }
  424. return "", nil
  425. }
  426. return "", nil
  427. }
  428. func (s *OpenAIGatewayService) validateCurrentGrokCredentialFailure(
  429. ctx context.Context,
  430. accountID int64,
  431. snapshot GrokCredentialMutationSnapshot,
  432. class grokCredentialFailureClass,
  433. ) (string, error) {
  434. if s == nil || s.accountRepo == nil || accountID <= 0 || ctx == nil || ctx.Err() != nil {
  435. if ctx != nil {
  436. return "", ctx.Err()
  437. }
  438. return "", context.Canceled
  439. }
  440. checkCtx, cancel := context.WithTimeout(ctx, time.Second)
  441. defer cancel()
  442. latest, err := s.accountRepo.GetByID(checkCtx, accountID)
  443. if err != nil {
  444. if errors.Is(err, ErrAccountNotFound) {
  445. return "", errOAuthRefreshAccountStateChanged
  446. }
  447. return "", fmt.Errorf("%w: %v", errOAuthRefreshAccountRereadFailed, err)
  448. }
  449. if latest == nil || !latest.IsGrokOAuth() || !latest.IsSchedulable() || s.isOpenAIAccountRuntimeBlocked(latest) {
  450. return "", errOAuthRefreshAccountStateChanged
  451. }
  452. latestSnapshot := grokCredentialMutationSnapshot(latest)
  453. if latestSnapshot.CredentialsJSON != snapshot.CredentialsJSON ||
  454. !grokCredentialProxyIDsEqual(latestSnapshot.ProxyID, snapshot.ProxyID) {
  455. if token, ok := s.grokCredentialConcurrentlyRefreshedToken(ctx, accountID, snapshot); ok {
  456. return token, nil
  457. }
  458. return "", errOAuthRefreshAccountStateChanged
  459. }
  460. // A configured proxy is external to the account-row CAS identity. Recheck
  461. // the hydrated proxy object so restoring a deleted row under the same ID
  462. // wins over a stale proxy-invalid failure.
  463. if class.reason == GrokCredentialReasonProxyInvalid {
  464. if latest.ProxyID == nil || latest.Proxy != nil {
  465. return "", errOAuthRefreshAccountStateChanged
  466. }
  467. } else if latest.ProxyID != nil && latest.Proxy == nil {
  468. return "", errOAuthRefreshAccountStateChanged
  469. }
  470. if class.reason == GrokCredentialReasonMissing {
  471. expiresAt := latest.GetCredentialAsTime("expires_at")
  472. credentialsStillMissing := strings.TrimSpace(latest.GetGrokAccessToken()) == "" ||
  473. strings.TrimSpace(latest.GetGrokRefreshToken()) == "" || expiresAt == nil || !time.Now().Before(*expiresAt)
  474. if !credentialsStillMissing {
  475. return "", errOAuthRefreshAccountStateChanged
  476. }
  477. }
  478. return "", nil
  479. }
  480. func (s *OpenAIGatewayService) grokCredentialMutationLock(accountID int64) *oauthRefreshLocalLock {
  481. actual, _ := s.grokCredentialMutationLocks.LoadOrStore(accountID, newOAuthRefreshLocalLock())
  482. mu, ok := actual.(*oauthRefreshLocalLock)
  483. if !ok {
  484. mu = newOAuthRefreshLocalLock()
  485. s.grokCredentialMutationLocks.Store(accountID, mu)
  486. }
  487. return mu
  488. }
  489. func (s *OpenAIGatewayService) grokCredentialMutationCommitted(accountID int64, class grokCredentialFailureClass, until time.Time) bool {
  490. if s == nil || s.accountRepo == nil || accountID <= 0 {
  491. return false
  492. }
  493. confirmCtx, cancel := context.WithTimeout(context.Background(), grokCredentialMutationConfirmWait)
  494. defer cancel()
  495. latest, err := s.accountRepo.GetByID(confirmCtx, accountID)
  496. if err != nil || latest == nil {
  497. return false
  498. }
  499. if class.permanent {
  500. return latest.Status == StatusError && !latest.Schedulable && latest.ErrorMessage == string(class.reason)
  501. }
  502. if class.transient {
  503. return latest.TempUnschedulableUntil != nil && !latest.TempUnschedulableUntil.Before(until) &&
  504. latest.TempUnschedulableReason == string(class.reason)
  505. }
  506. return false
  507. }
  508. func grokCredentialMutationSnapshot(account *Account) GrokCredentialMutationSnapshot {
  509. if account == nil {
  510. return GrokCredentialMutationSnapshot{}
  511. }
  512. credentialsJSON := "null"
  513. if encoded, err := json.Marshal(account.Credentials); err == nil {
  514. credentialsJSON = string(encoded)
  515. }
  516. snapshot := GrokCredentialMutationSnapshot{
  517. CredentialsJSON: credentialsJSON,
  518. AccessToken: strings.TrimSpace(account.GetGrokAccessToken()),
  519. RefreshToken: strings.TrimSpace(account.GetGrokRefreshToken()),
  520. TokenVersion: account.GetCredentialAsInt64("_token_version"),
  521. }
  522. if account.ProxyID != nil {
  523. proxyID := *account.ProxyID
  524. snapshot.ProxyID = &proxyID
  525. }
  526. return snapshot
  527. }
  528. func (s *OpenAIGatewayService) resolveGrokCredentialCASMiss(ctx context.Context, accountID int64, snapshot GrokCredentialMutationSnapshot) (string, error) {
  529. if ctx.Err() != nil {
  530. return "", ctx.Err()
  531. }
  532. if token, ok := s.grokCredentialConcurrentlyRefreshedToken(ctx, accountID, snapshot); ok {
  533. return token, nil
  534. }
  535. return "", errOAuthRefreshAccountStateChanged
  536. }
  537. func (s *OpenAIGatewayService) blockGrokCredentialRuntime(account *Account, until time.Time, reason string) func() {
  538. if s == nil || account == nil {
  539. return func() {}
  540. }
  541. mu := s.openAIAccountRuntimeBlockLock(account.ID)
  542. mu.Lock()
  543. before, hadBefore := s.openaiAccountRuntimeBlockUntil.Load(account.ID)
  544. installedGeneration, changed := s.blockAccountSchedulingLocked(account, until, reason)
  545. installed, installedOK := s.openaiAccountRuntimeBlockUntil.Load(account.ID)
  546. installedUntil, isTime := installed.(time.Time)
  547. mu.Unlock()
  548. if !changed || !installedOK || !isTime {
  549. return func() {}
  550. }
  551. if hadBefore {
  552. if beforeUntil, ok := before.(time.Time); ok && beforeUntil.Equal(installedUntil) {
  553. return func() {}
  554. }
  555. }
  556. return func() {
  557. mu.Lock()
  558. defer mu.Unlock()
  559. generation, ok := s.openaiAccountRuntimeBlockGeneration.Load(account.ID)
  560. if !ok || generation != installedGeneration {
  561. return
  562. }
  563. current, ok := s.openaiAccountRuntimeBlockUntil.Load(account.ID)
  564. currentUntil, isTime := current.(time.Time)
  565. if !ok || !isTime || !currentUntil.Equal(installedUntil) {
  566. return
  567. }
  568. if hadBefore {
  569. s.openaiAccountRuntimeBlockUntil.Store(account.ID, before)
  570. s.openaiAccountRuntimeBlockGeneration.Store(account.ID, s.openaiAccountRuntimeBlockSequence.Add(1))
  571. return
  572. }
  573. s.openaiAccountRuntimeBlockUntil.Delete(account.ID)
  574. s.openaiAccountRuntimeBlockGeneration.Store(account.ID, s.openaiAccountRuntimeBlockSequence.Add(1))
  575. }
  576. }
  577. func (s *OpenAIGatewayService) grokCredentialConcurrentlyRefreshedToken(ctx context.Context, accountID int64, baseline GrokCredentialMutationSnapshot) (string, bool) {
  578. if s == nil || s.accountRepo == nil || accountID <= 0 || ctx == nil || ctx.Err() != nil {
  579. return "", false
  580. }
  581. checkCtx, cancel := context.WithTimeout(ctx, time.Second)
  582. defer cancel()
  583. latest, err := s.accountRepo.GetByID(checkCtx, accountID)
  584. if err != nil || latest == nil {
  585. return "", false
  586. }
  587. latestSnapshot := grokCredentialMutationSnapshot(latest)
  588. if !grokCredentialProxyIDsEqual(latestSnapshot.ProxyID, baseline.ProxyID) ||
  589. latestSnapshot.CredentialsJSON == baseline.CredentialsJSON || !latest.IsSchedulable() ||
  590. (latest.ProxyID != nil && latest.Proxy == nil) || s.isOpenAIAccountRuntimeBlocked(latest) {
  591. return "", false
  592. }
  593. latestToken := strings.TrimSpace(latest.GetGrokAccessToken())
  594. if latestToken == "" || strings.TrimSpace(latest.GetGrokRefreshToken()) == "" {
  595. return "", false
  596. }
  597. expiresAt := latest.GetCredentialAsTime("expires_at")
  598. if expiresAt == nil || !time.Now().Before(*expiresAt) {
  599. return "", false
  600. }
  601. return latestToken, true
  602. }
  603. func grokCredentialProxyIDsEqual(left, right *int64) bool {
  604. if left == nil || right == nil {
  605. return left == nil && right == nil
  606. }
  607. return *left == *right
  608. }
  609. func (s *OpenAIGatewayService) newGrokCredentialFailover(c *gin.Context, account *Account, class grokCredentialFailureClass) error {
  610. if strings.TrimSpace(class.message) == "" {
  611. class.message = "Grok OAuth credentials are unavailable"
  612. }
  613. appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
  614. // Credential acquisition happens before the inference transport opens,
  615. // so the account binding is not evidence of an inference proxy route.
  616. ProxyID: nil,
  617. ProxyName: opsProxyNameUnknown,
  618. Platform: PlatformGrok,
  619. AccountID: account.ID,
  620. Stage: string(GatewayFailureStageAccountAuth),
  621. Scope: string(class.scope),
  622. Reason: string(class.reason),
  623. Kind: "credential_failover",
  624. Message: class.message,
  625. })
  626. return &UpstreamFailoverError{
  627. Stage: GatewayFailureStageAccountAuth,
  628. Scope: class.scope,
  629. Reason: class.reason,
  630. NextAccountAction: class.action,
  631. ClientStatusCode: http.StatusServiceUnavailable,
  632. ClientMessage: GrokCredentialUnavailableClientMessage,
  633. }
  634. }