Merge pull request #4221 from heathermhuang/codex/fix-grok-oauth-pool-health

fix(grok): refresh OAuth pools proactively
This commit is contained in:
Wesley Liddick
2026-07-15 16:07:09 +08:00
committed by GitHub
23 changed files with 5303 additions and 264 deletions
+2 -2
View File
@@ -201,7 +201,8 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
openAIOAuthHandler := admin.NewOpenAIOAuthHandler(openAIOAuthService, adminService, openAIQuotaService)
geminiOAuthHandler := admin.NewGeminiOAuthHandler(geminiOAuthService)
antigravityOAuthHandler := admin.NewAntigravityOAuthHandler(antigravityOAuthService)
grokOAuthHandler := admin.NewGrokOAuthHandler(grokOAuthService, adminService, grokQuotaService)
tokenRefreshService := service.ProvideTokenRefreshService(accountRepository, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, compositeTokenCacheInvalidator, schedulerCache, configConfig, tempUnschedCache, privacyClientFactory, proxyRepository, oAuthRefreshAPI, openAIGatewayService)
grokOAuthHandler := admin.NewGrokOAuthHandler(grokOAuthService, adminService, grokQuotaService, tokenRefreshService)
proxyHandler := admin.NewProxyHandler(adminService)
adminRedeemHandler := admin.NewRedeemHandler(adminService, redeemService)
promoHandler := admin.NewPromoHandler(promoService)
@@ -281,7 +282,6 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
opsAlertEvaluatorService := service.ProvideOpsAlertEvaluatorService(opsService, opsRepository, emailService, redisClient, configConfig, proxyRepository)
opsCleanupService := service.ProvideOpsCleanupService(opsRepository, db, redisClient, configConfig, channelMonitorService, settingRepository, opsService)
opsScheduledReportService := service.ProvideOpsScheduledReportService(opsService, userService, emailService, redisClient, configConfig)
tokenRefreshService := service.ProvideTokenRefreshService(accountRepository, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, compositeTokenCacheInvalidator, schedulerCache, configConfig, tempUnschedCache, privacyClientFactory, proxyRepository, oAuthRefreshAPI, openAIGatewayService)
accountExpiryService := service.ProvideAccountExpiryService(accountRepository)
proxyExpiryService := service.ProvideProxyExpiryService(proxyRepository)
subscriptionExpiryService := service.ProvideSubscriptionExpiryService(userSubscriptionRepository, settingRepository, notificationEmailService, leaderLockCache, db)
+18
View File
@@ -580,6 +580,18 @@ type TokenRefreshConfig struct {
MaxRetries int `mapstructure:"max_retries"`
// 重试退避基础时间(秒)
RetryBackoffSeconds int `mapstructure:"retry_backoff_seconds"`
// 每次从数据库读取的候选账号上限
CandidatePageSize int `mapstructure:"candidate_page_size"`
// 每个平台允许的并发刷新数
ProviderConcurrency int `mapstructure:"provider_concurrency"`
// 每个平台、每个进程允许的刷新请求速率
ProviderQPS int `mapstructure:"provider_qps"`
// 一个周期内连续临时失败达到此值后停止该平台
ProviderFailureThreshold int `mapstructure:"provider_failure_threshold"`
// 单次上游刷新尝试的超时(秒)
AttemptTimeoutSeconds int `mapstructure:"attempt_timeout_seconds"`
// 单个后台刷新周期的总超时(秒)
CycleTimeoutSeconds int `mapstructure:"cycle_timeout_seconds"`
}
type PricingConfig struct {
@@ -2112,6 +2124,12 @@ func setDefaults() {
viper.SetDefault("token_refresh.refresh_before_expiry_hours", 0.5) // 提前30分钟刷新(适配Google 1小时token
viper.SetDefault("token_refresh.max_retries", 3) // 最多重试3次
viper.SetDefault("token_refresh.retry_backoff_seconds", 2) // 重试退避基础2秒
viper.SetDefault("token_refresh.candidate_page_size", 200)
viper.SetDefault("token_refresh.provider_concurrency", 4)
viper.SetDefault("token_refresh.provider_qps", 2)
viper.SetDefault("token_refresh.provider_failure_threshold", 3)
viper.SetDefault("token_refresh.attempt_timeout_seconds", 15)
viper.SetDefault("token_refresh.cycle_timeout_seconds", 240)
// Gemini OAuth - configure via environment variables or config file
// GEMINI_OAUTH_CLIENT_ID and GEMINI_OAUTH_CLIENT_SECRET
@@ -73,7 +73,7 @@ func TestGrokSSOBatchImportKeepsCreatedAccountsWhenOneAutomaticProbeFails(t *tes
defer oauthService.Stop()
prober := newGrokImportProbeStub(3)
prober.failures[502] = infraerrors.New(502, "GROK_TEST_PROBE_FAILED", "sensitive-upstream-body")
handler := NewGrokOAuthHandler(oauthService, adminService, nil)
handler := NewGrokOAuthHandler(oauthService, adminService, nil, nil)
handler.importProber = prober
router := gin.New()
@@ -7,6 +7,7 @@ import (
"strconv"
"strings"
"sync"
"time"
"github.com/Wei-Shaw/sub2api/internal/handler/dto"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
@@ -23,18 +24,21 @@ type GrokOAuthHandler struct {
adminService service.AdminService
quotaService *service.GrokQuotaService
importProber grokUsageProber
reconciler service.GrokOAuthReconciler
}
func NewGrokOAuthHandler(
grokOAuthService *service.GrokOAuthService,
adminService service.AdminService,
quotaService *service.GrokQuotaService,
reconciler service.GrokOAuthReconciler,
) *GrokOAuthHandler {
return &GrokOAuthHandler{
grokOAuthService: grokOAuthService,
adminService: adminService,
quotaService: quotaService,
importProber: quotaService,
reconciler: reconciler,
}
}
@@ -160,6 +164,50 @@ func (h *GrokOAuthHandler) RefreshAccountToken(c *gin.Context) {
response.Success(c, dto.AccountFromService(updatedAccount))
}
type GrokOAuthReconcileRequest struct {
DryRun *bool `json:"dry_run"`
Apply bool `json:"apply"`
AfterID int64 `json:"after_id"`
Limit int `json:"limit"`
RefreshWindowSeconds int64 `json:"refresh_window_seconds"`
}
func (h *GrokOAuthHandler) ReconcileOAuthAccounts(c *gin.Context) {
var req GrokOAuthReconcileRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.BadRequest(c, "Invalid request")
return
}
dryRun := true
if req.DryRun != nil {
dryRun = *req.DryRun
}
if req.Apply == dryRun {
response.ErrorFrom(c, service.ErrGrokOAuthReconcileMode)
return
}
if req.RefreshWindowSeconds < 0 || req.RefreshWindowSeconds > int64((24*time.Hour)/time.Second) {
response.ErrorFrom(c, service.ErrGrokOAuthReconcileWindow)
return
}
if h.reconciler == nil {
response.InternalError(c, "Grok OAuth reconciliation service is unavailable")
return
}
result, err := h.reconciler.ReconcileGrokOAuth(c.Request.Context(), service.GrokOAuthReconcileInput{
DryRun: dryRun,
Apply: req.Apply,
AfterID: req.AfterID,
Limit: req.Limit,
RefreshWindow: time.Duration(req.RefreshWindowSeconds) * time.Second,
})
if err != nil {
response.ErrorFrom(c, err)
return
}
response.Success(c, result)
}
func (h *GrokOAuthHandler) CreateAccountFromOAuth(c *gin.Context) {
var req struct {
SessionID string `json:"session_id" binding:"required"`
@@ -47,6 +47,19 @@ type grokQuotaHandlerUpstream struct {
bodies [][]byte
}
type grokOAuthReconcilerStub struct {
input service.GrokOAuthReconcileInput
calls int
result *service.GrokOAuthReconcileResult
err error
}
func (s *grokOAuthReconcilerStub) ReconcileGrokOAuth(_ context.Context, input service.GrokOAuthReconcileInput) (*service.GrokOAuthReconcileResult, error) {
s.calls++
s.input = input
return s.result, s.err
}
func (u *grokQuotaHandlerUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
var body []byte
if req.Body != nil {
@@ -101,7 +114,7 @@ func TestGrokOAuthHandlerQueryQuotaProbesUpstream(t *testing.T) {
}}
upstream := &grokQuotaHandlerUpstream{}
quotaService := service.NewGrokQuotaService(repo, nil, service.NewGrokTokenProvider(repo, nil), upstream)
handler := NewGrokOAuthHandler(nil, nil, quotaService)
handler := NewGrokOAuthHandler(nil, nil, quotaService, nil)
router := gin.New()
router.GET("/api/v1/admin/grok/accounts/:id/quota", handler.QueryQuota)
@@ -139,7 +152,7 @@ func TestGrokOAuthHandlerResetQuotaReturnsUnsupported(t *testing.T) {
Type: service.AccountTypeOAuth,
}}
quotaService := service.NewGrokQuotaService(repo, nil, nil, nil)
handler := NewGrokOAuthHandler(nil, nil, quotaService)
handler := NewGrokOAuthHandler(nil, nil, quotaService, nil)
router := gin.New()
router.POST("/api/v1/admin/grok/accounts/:id/reset-quota", handler.ResetQuota)
@@ -157,7 +170,7 @@ func TestGrokOAuthHandlerRuntimeSanityDoesNotExposeSecrets(t *testing.T) {
t.Setenv(xai.EnvBaseURL, "http://127.0.0.1:8080/v1?access_token=secret")
t.Setenv(xai.EnvClientID, "client-secret-like-value")
handler := NewGrokOAuthHandler(nil, nil, nil)
handler := NewGrokOAuthHandler(nil, nil, nil, nil)
router := gin.New()
router.GET("/api/v1/admin/grok/runtime-sanity", handler.RuntimeSanity)
rec := httptest.NewRecorder()
@@ -219,3 +232,69 @@ func TestGrokSSOImportWorkerRecoversPanic(t *testing.T) {
require.Equal(t, 2, result.item.Index)
require.Contains(t, result.item.Error, "internal worker panic")
}
func TestGrokOAuthHandlerReconcileDefaultsToDryRun(t *testing.T) {
gin.SetMode(gin.TestMode)
reconciler := &grokOAuthReconcilerStub{result: &service.GrokOAuthReconcileResult{
DryRun: true,
Scanned: 2,
Actionable: 1,
WouldBlock: 1,
Items: []service.GrokOAuthReconcileItem{{AccountID: 42, Reason: service.GrokOAuthReconcileReasonMissingRefreshToken, Action: service.GrokOAuthReconcileActionBlock, Outcome: service.GrokOAuthReconcileOutcomePlanned}},
NextAfterID: 0,
}}
handler := NewGrokOAuthHandler(nil, nil, nil, reconciler)
router := gin.New()
router.POST("/api/v1/admin/grok/oauth/reconcile", handler.ReconcileOAuthAccounts)
rec := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/grok/oauth/reconcile", strings.NewReader(`{}`))
req.Header.Set("Content-Type", "application/json")
router.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
require.Equal(t, 1, reconciler.calls)
require.True(t, reconciler.input.DryRun)
require.False(t, reconciler.input.Apply)
require.Contains(t, rec.Body.String(), `"reason":"missing_refresh_token"`)
require.NotContains(t, rec.Body.String(), `"refresh_token":`)
require.NotContains(t, rec.Body.String(), `"access_token":`)
}
func TestGrokOAuthHandlerReconcileRequiresExplicitApply(t *testing.T) {
gin.SetMode(gin.TestMode)
reconciler := &grokOAuthReconcilerStub{}
handler := NewGrokOAuthHandler(nil, nil, nil, reconciler)
router := gin.New()
router.POST("/api/v1/admin/grok/oauth/reconcile", handler.ReconcileOAuthAccounts)
rec := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/grok/oauth/reconcile", strings.NewReader(`{"dry_run":false}`))
req.Header.Set("Content-Type", "application/json")
router.ServeHTTP(rec, req)
require.Equal(t, http.StatusBadRequest, rec.Code)
require.Zero(t, reconciler.calls)
require.NotContains(t, rec.Body.String(), "credentials")
}
func TestGrokOAuthHandlerReconcileExplicitApply(t *testing.T) {
gin.SetMode(gin.TestMode)
reconciler := &grokOAuthReconcilerStub{result: &service.GrokOAuthReconcileResult{DryRun: false, Refreshed: 1}}
handler := NewGrokOAuthHandler(nil, nil, nil, reconciler)
router := gin.New()
router.POST("/api/v1/admin/grok/oauth/reconcile", handler.ReconcileOAuthAccounts)
rec := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/grok/oauth/reconcile", strings.NewReader(`{"apply":true,"dry_run":false,"after_id":10,"limit":25,"refresh_window_seconds":3600}`))
req.Header.Set("Content-Type", "application/json")
router.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
require.Equal(t, 1, reconciler.calls)
require.True(t, reconciler.input.Apply)
require.False(t, reconciler.input.DryRun)
require.Equal(t, int64(10), reconciler.input.AfterID)
require.Equal(t, 25, reconciler.input.Limit)
require.Equal(t, time.Hour, reconciler.input.RefreshWindow)
}
+297 -12
View File
@@ -499,6 +499,9 @@ func (r *accountRepository) UpdateCredentials(ctx context.Context, id int64, cre
if err != nil {
return translatePersistenceError(err, service.ErrAccountNotFound, nil)
}
if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue credentials update failed: account=%d err=%v", id, err)
}
r.syncSchedulerAccountSnapshot(ctx, id)
return nil
}
@@ -783,29 +786,55 @@ func (r *accountRepository) ListActive(ctx context.Context) ([]service.Account,
return r.accountsToService(ctx, accounts)
}
func (r *accountRepository) ListOAuthRefreshCandidates(ctx context.Context) ([]service.Account, error) {
func (r *accountRepository) ListOAuthRefreshCandidatePage(ctx context.Context, options service.OAuthRefreshPageOptions) (*service.OAuthRefreshCandidatePage, error) {
if r.sql == nil {
return nil, errors.New("account repository SQL executor not configured")
}
if len(options.Platforms) == 0 {
return nil, errors.New("oauth refresh candidate platforms cannot be empty")
}
if options.Limit <= 0 || options.Limit > 1000 {
return nil, errors.New("oauth refresh candidate page limit must be between 1 and 1000")
}
// (cond) IS NOT TRUE 把 NULL 和 FALSE 都视为"可被刷新"。直接写
// NOT (a AND b) 在 PG 三值逻辑下会把 a 或 b 为 NULL 的行(即绝大多数
// 健康账号:temp_unschedulable_until=NULL)也排除,导致后台 token
// 刷新工作器漏掉所有正常账号 → access_token 到期后请求开始 401。
rows, err := r.sql.QueryContext(ctx, `
query := `
SELECT id
FROM accounts
WHERE deleted_at IS NULL
AND status = 'active'
AND type IN ('oauth', 'setup-token')
AND platform IN ('anthropic', 'openai', 'gemini', 'antigravity')
AND platform = ANY($1)
AND id > $2`
if options.ActiveOnly {
query += `
AND status = 'active'`
}
if options.IncludeSetupToken {
query += `
AND type IN ('oauth', 'setup-token')`
} else {
query += `
AND type = 'oauth'`
}
if options.RequireRefreshToken {
query += `
AND credentials ? 'refresh_token'
AND btrim(credentials->>'refresh_token') <> ''
AND btrim(credentials->>'refresh_token') <> ''`
}
if options.ExcludeRetryCooldown {
query += `
AND (
temp_unschedulable_until > NOW()
AND temp_unschedulable_reason LIKE 'token refresh retry exhausted:%'
) IS NOT TRUE
ORDER BY priority ASC, id ASC
`)
) IS NOT TRUE`
}
query += `
ORDER BY id ASC
LIMIT $3`
rows, err := r.sql.QueryContext(ctx, query, pq.Array(options.Platforms), options.AfterID, options.Limit)
if err != nil {
return nil, err
}
@@ -823,20 +852,33 @@ func (r *accountRepository) ListOAuthRefreshCandidates(ctx context.Context) ([]s
return nil, err
}
if len(ids) == 0 {
return []service.Account{}, nil
return &service.OAuthRefreshCandidatePage{Accounts: []service.Account{}}, nil
}
accounts, err := r.GetByIDs(ctx, ids)
if err != nil {
return nil, err
}
out := make([]service.Account, 0, len(accounts))
accountsByID := make(map[int64]*service.Account, len(accounts))
for _, account := range accounts {
if account != nil {
accountsByID[account.ID] = account
}
}
out := make([]service.Account, 0, len(accounts))
for _, id := range ids {
if account := accountsByID[id]; account != nil {
out = append(out, *account)
}
}
return out, nil
page := &service.OAuthRefreshCandidatePage{
Accounts: out,
HasMore: len(ids) == options.Limit,
}
if len(ids) > 0 {
page.NextAfterID = ids[len(ids)-1]
}
return page, nil
}
func (r *accountRepository) ListByPlatform(ctx context.Context, platform string) ([]service.Account, error) {
@@ -973,6 +1015,249 @@ func (r *accountRepository) SetGrokCredentialErrorIfMatch(
return true, nil
}
// SetGrokOAuthErrorIfCredentialsUnchanged atomically quarantines a structurally
// invalid Grok OAuth account only if it is still active and its complete JSONB
// credential document matches the state observed by reconciliation. Exact
// JSONB equality includes _token_version when present and prevents a concurrent
// reauthorization from being overwritten by a stale check-then-mutate path.
func (r *accountRepository) SetGrokOAuthErrorIfCredentialsUnchanged(
ctx context.Context,
id int64,
expectedCredentials map[string]any,
errorMsg string,
) (bool, error) {
if r == nil || r.sql == nil {
return false, errors.New("account repository SQL executor is not configured")
}
expectedJSON, err := json.Marshal(normalizeJSONMap(expectedCredentials))
if err != nil {
return false, err
}
result, err := r.sql.ExecContext(ctx, `
WITH updated AS (
UPDATE accounts AS a
SET status = $1,
error_message = $2,
schedulable = FALSE,
updated_at = NOW()
WHERE a.id = $3
AND a.deleted_at IS NULL
AND a.platform = $4
AND a.type = $5
AND a.status = $6
AND a.credentials = $7::jsonb
AND NULLIF(BTRIM(a.credentials->>'refresh_token'), '') IS NULL
RETURNING a.id
)
INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload)
SELECT $8, updated.id, NULL, NULL FROM updated
`,
service.StatusError,
errorMsg,
id,
service.PlatformGrok,
service.AccountTypeOAuth,
service.StatusActive,
string(expectedJSON),
service.SchedulerOutboxEventAccountChanged,
)
if err != nil {
return false, err
}
rowsAffected, err := result.RowsAffected()
if err != nil {
return false, err
}
if rowsAffected == 0 {
return false, nil
}
r.syncSchedulerAccountSnapshotDetached(ctx, id)
return true, nil
}
// UpdateGrokOAuthCredentialsIfUnchanged persists provider-issued replacement
// credentials only while the complete Grok OAuth credential document and
// proxy still match the fresh snapshot used by the upstream refresh call. The
// scheduler outbox insert is part of the same PostgreSQL statement, so a
// durable invalidation failure rolls the credential update back as well.
func (r *accountRepository) UpdateGrokOAuthCredentialsIfUnchanged(
ctx context.Context,
id int64,
expectedCredentials map[string]any,
expectedProxyID *int64,
credentials map[string]any,
) (bool, error) {
if r == nil || r.sql == nil {
return false, errors.New("account repository SQL executor is not configured")
}
expectedJSON, err := json.Marshal(normalizeJSONMap(expectedCredentials))
if err != nil {
return false, err
}
credentialsJSON, err := json.Marshal(normalizeJSONMap(credentials))
if err != nil {
return false, err
}
result, err := r.sql.ExecContext(ctx, `
WITH updated AS (
UPDATE accounts AS a
SET credentials = $1::jsonb,
updated_at = NOW()
WHERE a.id = $2
AND a.deleted_at IS NULL
AND a.platform = $3
AND a.type = $4
AND a.credentials = $5::jsonb
AND a.proxy_id IS NOT DISTINCT FROM $6
RETURNING a.id
)
INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload)
SELECT $7, updated.id, NULL, NULL FROM updated
`,
string(credentialsJSON),
id,
service.PlatformGrok,
service.AccountTypeOAuth,
string(expectedJSON),
expectedProxyID,
service.SchedulerOutboxEventAccountChanged,
)
if err != nil {
return false, err
}
rowsAffected, err := result.RowsAffected()
if err != nil {
return false, err
}
if rowsAffected == 0 {
return false, nil
}
r.syncSchedulerAccountSnapshotDetached(ctx, id)
return true, nil
}
// SetGrokOAuthRefreshErrorIfCredentialsUnchanged is the background-refresh
// counterpart to reconciliation's stricter missing-refresh-token mutation. It
// matches the complete credential document used by the failed upstream attempt
// but deliberately does not require the refresh token to be absent.
func (r *accountRepository) SetGrokOAuthRefreshErrorIfCredentialsUnchanged(
ctx context.Context,
id int64,
expectedCredentials map[string]any,
expectedProxyID *int64,
errorMsg string,
) (bool, error) {
if r == nil || r.sql == nil {
return false, errors.New("account repository SQL executor is not configured")
}
expectedJSON, err := json.Marshal(normalizeJSONMap(expectedCredentials))
if err != nil {
return false, err
}
result, err := r.sql.ExecContext(ctx, `
WITH updated AS (
UPDATE accounts AS a
SET status = $1,
error_message = $2,
schedulable = FALSE,
updated_at = NOW()
WHERE a.id = $3
AND a.deleted_at IS NULL
AND a.platform = $4
AND a.type = $5
AND a.status = $6
AND a.credentials = $7::jsonb
AND a.proxy_id IS NOT DISTINCT FROM $8
RETURNING a.id
)
INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload)
SELECT $9, updated.id, NULL, NULL FROM updated
`,
service.StatusError,
errorMsg,
id,
service.PlatformGrok,
service.AccountTypeOAuth,
service.StatusActive,
string(expectedJSON),
expectedProxyID,
service.SchedulerOutboxEventAccountChanged,
)
if err != nil {
return false, err
}
rowsAffected, err := result.RowsAffected()
if err != nil {
return false, err
}
if rowsAffected == 0 {
return false, nil
}
r.syncSchedulerAccountSnapshotDetached(ctx, id)
return true, nil
}
// SetGrokOAuthRefreshTempUnschedulableIfCredentialsUnchanged applies a bounded
// transient refresh quarantine only while the active Grok OAuth credential
// document still matches the exact upstream attempt.
func (r *accountRepository) SetGrokOAuthRefreshTempUnschedulableIfCredentialsUnchanged(
ctx context.Context,
id int64,
expectedCredentials map[string]any,
expectedProxyID *int64,
until time.Time,
reason string,
) (bool, error) {
if r == nil || r.sql == nil {
return false, errors.New("account repository SQL executor is not configured")
}
expectedJSON, err := json.Marshal(normalizeJSONMap(expectedCredentials))
if err != nil {
return false, err
}
result, err := r.sql.ExecContext(ctx, `
WITH updated AS (
UPDATE accounts AS a
SET temp_unschedulable_until = $1,
temp_unschedulable_reason = $2,
updated_at = NOW()
WHERE a.id = $3
AND a.deleted_at IS NULL
AND a.platform = $4
AND a.type = $5
AND a.status = $6
AND a.credentials = $7::jsonb
AND a.proxy_id IS NOT DISTINCT FROM $8
AND (a.temp_unschedulable_until IS NULL OR a.temp_unschedulable_until < $1)
RETURNING a.id
)
INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload)
SELECT $9, updated.id, NULL, NULL FROM updated
`,
until,
reason,
id,
service.PlatformGrok,
service.AccountTypeOAuth,
service.StatusActive,
string(expectedJSON),
expectedProxyID,
service.SchedulerOutboxEventAccountChanged,
)
if err != nil {
return false, err
}
rowsAffected, err := result.RowsAffected()
if err != nil {
return false, err
}
if rowsAffected == 0 {
return false, nil
}
r.syncSchedulerAccountSnapshotDetached(ctx, id)
return true, nil
}
// syncSchedulerAccountSnapshot 在账号状态变更时主动同步快照到调度器缓存。
// 当账号被设置为错误、禁用、不可调度或临时不可调度时调用,
// 确保调度器和粘性会话逻辑能及时感知账号的最新状态,避免继续使用不可用账号。
@@ -4,6 +4,8 @@ package repository
import (
"context"
"database/sql"
"strings"
"testing"
"time"
@@ -11,6 +13,7 @@ import (
"github.com/Wei-Shaw/sub2api/ent/accountgroup"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/stretchr/testify/require"
"github.com/stretchr/testify/suite"
)
@@ -25,6 +28,7 @@ type schedulerCacheRecorder struct {
setAccounts []*service.Account
deleteIDs []int64
accounts map[int64]*service.Account
setCtxErr error
}
func (s *schedulerCacheRecorder) GetSnapshot(ctx context.Context, bucket service.SchedulerBucket) ([]*service.Account, bool, error) {
@@ -63,6 +67,7 @@ func (s *schedulerCacheRecorder) GetAccount(ctx context.Context, accountID int64
}
func (s *schedulerCacheRecorder) SetAccount(ctx context.Context, account *service.Account) error {
s.setCtxErr = ctx.Err()
s.setAccounts = append(s.setAccounts, account)
if s.accounts == nil {
s.accounts = make(map[int64]*service.Account)
@@ -73,6 +78,31 @@ func (s *schedulerCacheRecorder) SetAccount(ctx context.Context, account *servic
return nil
}
type failAtomicSchedulerOutboxSQLExecutor struct {
sqlExecutor
}
func (e *failAtomicSchedulerOutboxSQLExecutor) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
if strings.Contains(query, "WITH updated AS") && strings.Contains(query, "INSERT INTO scheduler_outbox") && len(args) > 0 {
args = append([]any(nil), args...)
args[len(args)-1] = nil // event_type is NOT NULL; the whole statement must roll back.
}
return e.sqlExecutor.ExecContext(ctx, query, args...)
}
type cancelAfterAtomicMutationSQLExecutor struct {
sqlExecutor
cancel context.CancelFunc
}
func (e *cancelAfterAtomicMutationSQLExecutor) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
result, err := e.sqlExecutor.ExecContext(ctx, query, args...)
if err == nil && strings.Contains(query, "WITH updated AS") && strings.Contains(query, "INSERT INTO scheduler_outbox") {
e.cancel()
}
return result, err
}
func (s *schedulerCacheRecorder) DeleteAccount(ctx context.Context, accountID int64) error {
s.deleteIDs = append(s.deleteIDs, accountID)
if s.accounts != nil {
@@ -200,6 +230,34 @@ func (s *AccountRepoSuite) TestUpdate_SyncSchedulerSnapshotOnCredentialsChange()
s.Require().Equal("gpt-5.2", mapping["gpt-5"])
}
func (s *AccountRepoSuite) TestUpdateCredentials_SyncsSnapshotAndDurableOutbox() {
account := mustCreateAccount(s.T(), s.client, &service.Account{
Name: "sync-refresh-credentials",
Status: service.StatusActive,
Schedulable: true,
Credentials: map[string]any{"access_token": "old-token"},
})
cacheRecorder := &schedulerCacheRecorder{}
s.repo.schedulerCache = cacheRecorder
_, err := s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
s.Require().NoError(err)
s.Require().NoError(s.repo.UpdateCredentials(s.ctx, account.ID, map[string]any{"access_token": "new-token"}))
s.Require().Len(cacheRecorder.setAccounts, 1)
s.Require().Equal("new-token", cacheRecorder.setAccounts[0].GetCredential("access_token"))
var outboxCount int
err = scanSingleRow(
s.ctx,
s.repo.sql,
"SELECT COUNT(*) FROM scheduler_outbox WHERE event_type = $1 AND account_id = $2",
[]any{service.SchedulerOutboxEventAccountChanged, account.ID},
&outboxCount,
)
s.Require().NoError(err)
s.Require().Equal(1, outboxCount)
}
func (s *AccountRepoSuite) TestDelete() {
account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "to-delete"})
@@ -256,6 +314,88 @@ func (s *AccountRepoSuite) TestList() {
s.Require().Equal(int64(2), page.Total)
}
func (s *AccountRepoSuite) TestListOAuthRefreshCandidatePage_GrokCursorAndExclusions() {
now := time.Now().UTC()
valid1 := mustCreateAccount(s.T(), s.client, &service.Account{
Name: "grok-oauth-page-1",
Platform: service.PlatformGrok,
Type: service.AccountTypeOAuth,
Status: service.StatusActive,
Credentials: map[string]any{
"access_token": "access-1",
"refresh_token": "refresh-1",
"expires_at": now.Add(30 * time.Minute).Format(time.RFC3339),
},
})
mustCreateAccount(s.T(), s.client, &service.Account{
Name: "grok-api-key-excluded",
Platform: service.PlatformGrok,
Type: service.AccountTypeAPIKey,
Status: service.StatusActive,
Credentials: map[string]any{
"api_key": "api-key",
"refresh_token": "must-not-make-api-key-eligible",
},
})
valid2 := mustCreateAccount(s.T(), s.client, &service.Account{
Name: "grok-oauth-page-2",
Platform: service.PlatformGrok,
Type: service.AccountTypeOAuth,
Status: service.StatusActive,
Credentials: map[string]any{"refresh_token": "refresh-2"},
})
mustCreateAccount(s.T(), s.client, &service.Account{
Name: "grok-oauth-blank-refresh-excluded",
Platform: service.PlatformGrok,
Type: service.AccountTypeOAuth,
Status: service.StatusActive,
Credentials: map[string]any{"refresh_token": " "},
})
valid3 := mustCreateAccount(s.T(), s.client, &service.Account{
Name: "grok-oauth-page-3",
Platform: service.PlatformGrok,
Type: service.AccountTypeOAuth,
Status: service.StatusActive,
Credentials: map[string]any{"refresh_token": "refresh-3"},
})
cooldown := mustCreateAccount(s.T(), s.client, &service.Account{
Name: "grok-oauth-retry-cooldown-excluded",
Platform: service.PlatformGrok,
Type: service.AccountTypeOAuth,
Status: service.StatusActive,
Credentials: map[string]any{"refresh_token": "refresh-cooldown"},
})
s.Require().NoError(s.repo.SetTempUnschedulable(s.ctx, cooldown.ID, now.Add(10*time.Minute), "token refresh retry exhausted: timeout"))
mustCreateAccount(s.T(), s.client, &service.Account{
Name: "openai-oauth-excluded",
Platform: service.PlatformOpenAI,
Type: service.AccountTypeOAuth,
Status: service.StatusActive,
Credentials: map[string]any{"refresh_token": "refresh-openai"},
})
options := service.OAuthRefreshPageOptions{
Platforms: []string{service.PlatformGrok},
Limit: 2,
ActiveOnly: true,
RequireRefreshToken: true,
ExcludeRetryCooldown: true,
}
firstPage, err := s.repo.ListOAuthRefreshCandidatePage(s.ctx, options)
s.Require().NoError(err)
first := firstPage.Accounts
s.Require().Len(first, 2)
s.Require().Equal([]int64{valid1.ID, valid2.ID}, []int64{first[0].ID, first[1].ID})
options.AfterID = first[len(first)-1].ID
secondPage, err := s.repo.ListOAuthRefreshCandidatePage(s.ctx, options)
s.Require().NoError(err)
second := secondPage.Accounts
s.Require().Len(second, 1)
s.Require().Equal(valid3.ID, second[0].ID)
s.Require().NotContains([]int64{first[0].ID, first[1].ID}, second[0].ID)
}
func (s *AccountRepoSuite) TestListWithFilters() {
tests := []struct {
name string
@@ -918,6 +1058,10 @@ func (s *AccountRepoSuite) TestUpdateLastUsed() {
func (s *AccountRepoSuite) TestSetError() {
account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-err", Status: service.StatusActive, Schedulable: true})
cacheRecorder := &schedulerCacheRecorder{}
s.repo.schedulerCache = cacheRecorder
_, err := s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
s.Require().NoError(err)
s.Require().NoError(s.repo.SetError(s.ctx, account.ID, "something went wrong"))
@@ -926,6 +1070,296 @@ func (s *AccountRepoSuite) TestSetError() {
s.Require().Equal(service.StatusError, got.Status)
s.Require().Equal("something went wrong", got.ErrorMessage)
s.Require().False(got.Schedulable)
s.Require().Len(cacheRecorder.setAccounts, 1)
s.Require().Equal(account.ID, cacheRecorder.setAccounts[0].ID)
s.Require().Equal(service.StatusError, cacheRecorder.setAccounts[0].Status)
s.Require().False(cacheRecorder.setAccounts[0].Schedulable)
var outboxCount int
err = scanSingleRow(
s.ctx,
s.repo.sql,
"SELECT COUNT(*) FROM scheduler_outbox WHERE event_type = $1 AND account_id = $2",
[]any{service.SchedulerOutboxEventAccountChanged, account.ID},
&outboxCount,
)
s.Require().NoError(err)
s.Require().Equal(1, outboxCount)
}
func (s *AccountRepoSuite) TestSetGrokOAuthErrorIfCredentialsUnchanged_AppliesAndSyncsSchedulerState() {
account := mustCreateAccount(s.T(), s.client, &service.Account{
Name: "grok-conditional-error-applied",
Platform: service.PlatformGrok,
Type: service.AccountTypeOAuth,
Status: service.StatusActive,
Schedulable: true,
Credentials: map[string]any{"access_token": "observed", "_token_version": int64(7)},
})
observed, err := s.repo.GetByID(s.ctx, account.ID)
s.Require().NoError(err)
cacheRecorder := &schedulerCacheRecorder{}
s.repo.schedulerCache = cacheRecorder
_, err = s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
s.Require().NoError(err)
applied, err := s.repo.SetGrokOAuthErrorIfCredentialsUnchanged(
s.ctx,
account.ID,
observed.Credentials,
"missing refresh token",
)
s.Require().NoError(err)
s.Require().True(applied)
got, err := s.repo.GetByID(s.ctx, account.ID)
s.Require().NoError(err)
s.Require().Equal(service.StatusError, got.Status)
s.Require().False(got.Schedulable)
s.Require().Equal("missing refresh token", got.ErrorMessage)
s.Require().Len(cacheRecorder.setAccounts, 1)
s.Require().Equal(service.StatusError, cacheRecorder.setAccounts[0].Status)
var outboxCount int
err = scanSingleRow(
s.ctx,
s.repo.sql,
"SELECT COUNT(*) FROM scheduler_outbox WHERE event_type = $1 AND account_id = $2",
[]any{service.SchedulerOutboxEventAccountChanged, account.ID},
&outboxCount,
)
s.Require().NoError(err)
s.Require().Equal(1, outboxCount)
}
func (s *AccountRepoSuite) TestSetGrokOAuthErrorIfCredentialsUnchanged_SkipsConcurrentReauthorization() {
account := mustCreateAccount(s.T(), s.client, &service.Account{
Name: "grok-conditional-error-reauthorized",
Platform: service.PlatformGrok,
Type: service.AccountTypeOAuth,
Status: service.StatusActive,
Schedulable: true,
Credentials: map[string]any{"access_token": "observed", "_token_version": int64(7)},
})
observed, err := s.repo.GetByID(s.ctx, account.ID)
s.Require().NoError(err)
s.Require().NoError(s.repo.UpdateCredentials(s.ctx, account.ID, map[string]any{
"access_token": "fresh-access",
"refresh_token": "fresh-refresh",
"expires_at": time.Now().UTC().Add(4 * time.Hour).Format(time.RFC3339),
"_token_version": int64(8),
}))
cacheRecorder := &schedulerCacheRecorder{}
s.repo.schedulerCache = cacheRecorder
_, err = s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
s.Require().NoError(err)
applied, err := s.repo.SetGrokOAuthErrorIfCredentialsUnchanged(
s.ctx,
account.ID,
observed.Credentials,
"stale reconciliation",
)
s.Require().NoError(err)
s.Require().False(applied)
got, err := s.repo.GetByID(s.ctx, account.ID)
s.Require().NoError(err)
s.Require().Equal(service.StatusActive, got.Status)
s.Require().True(got.Schedulable)
s.Require().Equal("fresh-refresh", got.GetGrokRefreshToken())
s.Require().Empty(cacheRecorder.setAccounts, "a lost compare-and-set race must not rewrite the scheduler snapshot")
var outboxCount int
err = scanSingleRow(
s.ctx,
s.repo.sql,
"SELECT COUNT(*) FROM scheduler_outbox WHERE event_type = $1 AND account_id = $2",
[]any{service.SchedulerOutboxEventAccountChanged, account.ID},
&outboxCount,
)
s.Require().NoError(err)
s.Require().Zero(outboxCount, "a lost compare-and-set race must not enqueue a stale account change")
}
func (s *AccountRepoSuite) TestUpdateGrokOAuthCredentialsIfUnchanged_AppliesAndPublishesSchedulerState() {
account := mustCreateAccount(s.T(), s.client, &service.Account{
Name: "grok-refresh-success-cas-applied",
Platform: service.PlatformGrok,
Type: service.AccountTypeOAuth,
Status: service.StatusActive,
Schedulable: true,
Credentials: map[string]any{
"access_token": "attempted-access",
"refresh_token": "attempted-refresh",
"_token_version": int64(10),
},
})
observed, err := s.repo.GetByID(s.ctx, account.ID)
s.Require().NoError(err)
cacheRecorder := &schedulerCacheRecorder{}
s.repo.schedulerCache = cacheRecorder
_, err = s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
s.Require().NoError(err)
applied, err := s.repo.UpdateGrokOAuthCredentialsIfUnchanged(
s.ctx,
account.ID,
observed.Credentials,
observed.ProxyID,
map[string]any{
"access_token": "rotated-access",
"refresh_token": "rotated-refresh",
"_token_version": int64(11),
},
)
s.Require().NoError(err)
s.Require().True(applied)
got, err := s.repo.GetByID(s.ctx, account.ID)
s.Require().NoError(err)
s.Require().Equal("rotated-refresh", got.GetGrokRefreshToken())
s.Require().Len(cacheRecorder.setAccounts, 1)
s.Require().Equal("rotated-refresh", cacheRecorder.setAccounts[0].GetGrokRefreshToken())
s.Require().NoError(cacheRecorder.setCtxErr)
var outboxCount int
err = scanSingleRow(
s.ctx,
s.repo.sql,
"SELECT COUNT(*) FROM scheduler_outbox WHERE event_type = $1 AND account_id = $2",
[]any{service.SchedulerOutboxEventAccountChanged, account.ID},
&outboxCount,
)
s.Require().NoError(err)
s.Require().Equal(1, outboxCount)
}
func (s *AccountRepoSuite) TestUpdateGrokOAuthCredentialsIfUnchanged_SkipsConcurrentReauthorization() {
account := mustCreateAccount(s.T(), s.client, &service.Account{
Name: "grok-refresh-success-cas-reauthorized",
Platform: service.PlatformGrok,
Type: service.AccountTypeOAuth,
Status: service.StatusActive,
Schedulable: true,
Credentials: map[string]any{
"access_token": "attempted-access",
"refresh_token": "attempted-refresh",
"_token_version": int64(20),
},
})
observed, err := s.repo.GetByID(s.ctx, account.ID)
s.Require().NoError(err)
s.Require().NoError(s.repo.UpdateCredentials(s.ctx, account.ID, map[string]any{
"access_token": "reauthorized-access",
"refresh_token": "reauthorized-refresh",
"_token_version": int64(21),
}))
cacheRecorder := &schedulerCacheRecorder{}
s.repo.schedulerCache = cacheRecorder
_, err = s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
s.Require().NoError(err)
applied, err := s.repo.UpdateGrokOAuthCredentialsIfUnchanged(
s.ctx,
account.ID,
observed.Credentials,
observed.ProxyID,
map[string]any{
"access_token": "provider-access",
"refresh_token": "provider-refresh",
"_token_version": int64(22),
},
)
s.Require().NoError(err)
s.Require().False(applied)
got, err := s.repo.GetByID(s.ctx, account.ID)
s.Require().NoError(err)
s.Require().Equal("reauthorized-refresh", got.GetGrokRefreshToken())
s.Require().Empty(cacheRecorder.setAccounts)
var outboxCount int
err = scanSingleRow(
s.ctx,
s.repo.sql,
"SELECT COUNT(*) FROM scheduler_outbox WHERE event_type = $1 AND account_id = $2",
[]any{service.SchedulerOutboxEventAccountChanged, account.ID},
&outboxCount,
)
s.Require().NoError(err)
s.Require().Zero(outboxCount)
}
func (s *AccountRepoSuite) TestGrokOAuthConditionalMutation_DetachesBoundedSnapshotSync() {
account := mustCreateAccount(s.T(), s.client, &service.Account{
Name: "grok-conditional-detached-sync",
Platform: service.PlatformGrok,
Type: service.AccountTypeOAuth,
Status: service.StatusActive,
Schedulable: true,
Credentials: map[string]any{"access_token": "observed"},
})
observed, err := s.repo.GetByID(s.ctx, account.ID)
s.Require().NoError(err)
ctx, cancel := context.WithCancel(context.Background())
cacheRecorder := &schedulerCacheRecorder{}
repo := newAccountRepositoryWithSQL(s.client, &cancelAfterAtomicMutationSQLExecutor{
sqlExecutor: s.repo.sql,
cancel: cancel,
}, cacheRecorder)
applied, err := repo.SetGrokOAuthErrorIfCredentialsUnchanged(
ctx,
account.ID,
observed.Credentials,
"missing refresh token",
)
s.Require().NoError(err)
s.Require().True(applied)
s.Require().ErrorIs(ctx.Err(), context.Canceled)
s.Require().Len(cacheRecorder.setAccounts, 1)
s.Require().NoError(cacheRecorder.setCtxErr, "immediate scheduler propagation must use a bounded detached context")
}
func TestGrokOAuthConditionalMutationRollsBackWhenOutboxInsertFails(t *testing.T) {
client := testEntClient(t)
account := mustCreateAccount(t, client, &service.Account{
Name: "grok-conditional-atomic-outbox-failure",
Platform: service.PlatformGrok,
Type: service.AccountTypeOAuth,
Status: service.StatusActive,
Schedulable: true,
Credentials: map[string]any{"access_token": "observed"},
})
t.Cleanup(func() {
_, _ = integrationDB.ExecContext(context.Background(), "DELETE FROM scheduler_outbox WHERE account_id = $1", account.ID)
_ = client.Account.DeleteOneID(account.ID).Exec(context.Background())
})
repo := newAccountRepositoryWithSQL(client, &failAtomicSchedulerOutboxSQLExecutor{sqlExecutor: integrationDB}, nil)
applied, err := repo.SetGrokOAuthErrorIfCredentialsUnchanged(
context.Background(),
account.ID,
account.Credentials,
"missing refresh token",
)
require.Error(t, err)
require.False(t, applied)
got, readErr := repo.GetByID(context.Background(), account.ID)
require.NoError(t, readErr)
require.Equal(t, service.StatusActive, got.Status)
require.True(t, got.Schedulable)
require.Empty(t, got.ErrorMessage)
var outboxCount int
require.NoError(t, integrationDB.QueryRowContext(
context.Background(),
"SELECT COUNT(*) FROM scheduler_outbox WHERE account_id = $1",
account.ID,
).Scan(&outboxCount))
require.Zero(t, outboxCount)
}
func (s *AccountRepoSuite) TestUpdateErrorStatusUnschedulesAccount() {
@@ -3,6 +3,7 @@ package repository
import (
"context"
"database/sql"
"database/sql/driver"
"regexp"
"strings"
"testing"
@@ -122,28 +123,165 @@ func TestAccountRepository_GrokCredentialCommitCarriesOutboxAcrossCallerCancella
}
}
func TestAccountRepository_ListOAuthRefreshCandidates_SQLFilter(t *testing.T) {
func TestAccountRepository_SetGrokOAuthErrorIfCredentialsUnchanged_RequiresActiveExactCredentialMatch(t *testing.T) {
exec := &recordingSQLExecutor{result: rowsAffectedResult(0)}
repo := newAccountRepositoryWithSQL(nil, exec, nil)
applied, err := repo.SetGrokOAuthErrorIfCredentialsUnchanged(
context.Background(),
42,
map[string]any{"access_token": "observed", "_token_version": int64(7)},
"missing refresh token",
)
require.NoError(t, err)
require.False(t, applied)
require.Len(t, exec.execQueries, 1, "the account mutation and conditional outbox insert must be one statement")
normalized := normalizeSQLWhitespace(exec.execQueries[0])
require.Contains(t, normalized, "WITH updated AS")
require.Contains(t, normalized, "INSERT INTO scheduler_outbox")
require.Contains(t, normalized, "FROM updated")
require.Contains(t, normalized, "platform = $4")
require.Contains(t, normalized, "type = $5")
require.Contains(t, normalized, "status = $6")
require.Contains(t, normalized, "credentials = $7::jsonb")
require.Contains(t, normalized, "NULLIF(BTRIM(a.credentials->>'refresh_token'), '') IS NULL")
require.Len(t, exec.execArgs, 1)
require.Equal(t, service.StatusActive, exec.execArgs[0][5])
require.Contains(t, exec.execArgs[0][6], `"_token_version":7`)
}
func TestAccountRepository_SetGrokOAuthErrorIfCredentialsUnchanged_AppliedWritesOutbox(t *testing.T) {
exec := &recordingSQLExecutor{result: rowsAffectedResult(1)}
repo := newAccountRepositoryWithSQL(nil, exec, nil)
applied, err := repo.SetGrokOAuthErrorIfCredentialsUnchanged(
context.Background(),
42,
map[string]any{"access_token": "observed"},
"missing refresh token",
)
require.NoError(t, err)
require.True(t, applied)
require.Len(t, exec.execQueries, 1)
normalized := normalizeSQLWhitespace(exec.execQueries[0])
require.Contains(t, normalized, "WITH updated AS")
require.Contains(t, normalized, "INSERT INTO scheduler_outbox")
require.Contains(t, normalized, "SELECT $8, updated.id, NULL, NULL FROM updated")
}
func TestAccountRepository_SetGrokOAuthRefreshErrorIfCredentialsUnchanged_UsesAttemptCredentialsAndProxy(t *testing.T) {
exec := &recordingSQLExecutor{result: rowsAffectedResult(0)}
repo := newAccountRepositoryWithSQL(nil, exec, nil)
proxyID := int64(17)
applied, err := repo.SetGrokOAuthRefreshErrorIfCredentialsUnchanged(
context.Background(),
42,
map[string]any{"refresh_token": "attempted", "_token_version": int64(7)},
&proxyID,
"revoked",
)
require.NoError(t, err)
require.False(t, applied)
require.Len(t, exec.execQueries, 1)
normalized := normalizeSQLWhitespace(exec.execQueries[0])
require.Contains(t, normalized, "credentials = $7::jsonb")
require.Contains(t, normalized, "proxy_id IS NOT DISTINCT FROM $8")
require.NotContains(t, normalized, "credentials->>'refresh_token'",
"background invalid_grant CAS must accept the attempted refresh token; only reconciliation requires it missing")
require.Equal(t, &proxyID, exec.execArgs[0][7])
require.Contains(t, normalized, "INSERT INTO scheduler_outbox")
require.Len(t, exec.execArgs[0], 9)
}
func TestAccountRepository_SetGrokOAuthRefreshTempUnschedulableIfCredentialsUnchanged_UsesAttemptCredentialsAndProxy(t *testing.T) {
exec := &recordingSQLExecutor{result: rowsAffectedResult(0)}
repo := newAccountRepositoryWithSQL(nil, exec, nil)
proxyID := int64(19)
applied, err := repo.SetGrokOAuthRefreshTempUnschedulableIfCredentialsUnchanged(
context.Background(),
42,
map[string]any{"refresh_token": "attempted", "_token_version": int64(8)},
&proxyID,
time.Now().Add(10*time.Minute),
"retry exhausted",
)
require.NoError(t, err)
require.False(t, applied)
require.Len(t, exec.execQueries, 1)
normalized := normalizeSQLWhitespace(exec.execQueries[0])
require.Contains(t, normalized, "credentials = $7::jsonb")
require.Contains(t, normalized, "proxy_id IS NOT DISTINCT FROM $8")
require.Contains(t, normalized, "a.temp_unschedulable_until IS NULL OR a.temp_unschedulable_until < $1")
require.Len(t, exec.execArgs[0], 9)
require.Equal(t, &proxyID, exec.execArgs[0][7])
require.Contains(t, normalized, "INSERT INTO scheduler_outbox")
}
func TestAccountRepository_UpdateGrokOAuthCredentialsIfUnchanged_UsesExactAttemptStateAndAtomicOutbox(t *testing.T) {
exec := &recordingSQLExecutor{result: rowsAffectedResult(1)}
repo := newAccountRepositoryWithSQL(nil, exec, nil)
proxyID := int64(29)
applied, err := repo.UpdateGrokOAuthCredentialsIfUnchanged(
context.Background(),
42,
map[string]any{"refresh_token": "attempted", "_token_version": int64(9)},
&proxyID,
map[string]any{"refresh_token": "rotated", "_token_version": int64(10)},
)
require.NoError(t, err)
require.True(t, applied)
require.Len(t, exec.execQueries, 1)
normalized := normalizeSQLWhitespace(exec.execQueries[0])
require.Contains(t, normalized, "WITH updated AS")
require.Contains(t, normalized, "credentials = $1::jsonb")
require.Contains(t, normalized, "credentials = $5::jsonb")
require.Contains(t, normalized, "proxy_id IS NOT DISTINCT FROM $6")
require.Contains(t, normalized, "INSERT INTO scheduler_outbox")
require.Len(t, exec.execArgs[0], 7)
require.Equal(t, &proxyID, exec.execArgs[0][5])
}
func TestAccountRepository_ListOAuthRefreshCandidatePage_SQLFilter(t *testing.T) {
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
require.NoError(t, err)
defer func() { _ = db.Close() }()
var capturedSQL string
var capturedArgs []any
mock.ExpectQuery("SELECT id").
WillReturnRows(sqlmock.NewRows([]string{"id"})).
WillDelayFor(0)
repo := newAccountRepositoryWithSQL(nil, captureQuerySQL{db: db, captured: &capturedSQL}, nil)
repo := newAccountRepositoryWithSQL(nil, captureQuerySQL{db: db, captured: &capturedSQL, args: &capturedArgs}, nil)
accounts, err := repo.ListOAuthRefreshCandidates(context.Background())
page, err := repo.ListOAuthRefreshCandidatePage(context.Background(), service.OAuthRefreshPageOptions{
Platforms: []string{service.PlatformAnthropic, service.PlatformOpenAI, service.PlatformGemini, service.PlatformAntigravity, service.PlatformGrok},
AfterID: 100,
Limit: 200,
ActiveOnly: true,
IncludeSetupToken: true,
RequireRefreshToken: true,
ExcludeRetryCooldown: true,
})
require.NoError(t, err)
require.Empty(t, accounts)
require.Empty(t, page.Accounts)
normalized := normalizeSQLWhitespace(capturedSQL)
require.Contains(t, normalized, "deleted_at IS NULL")
require.Contains(t, normalized, "status = 'active'")
// setup-token 的 access_token 同为 8h 短期令牌,必须与 oauth 一起纳入后台刷新候选
require.Contains(t, normalized, "type IN ('oauth', 'setup-token')")
require.Contains(t, normalized, "platform IN ('anthropic', 'openai', 'gemini', 'antigravity')")
require.Contains(t, normalized, "platform = ANY($1)")
require.NotContains(t, normalized, "platform IN ('anthropic'",
"candidate platforms must come from the refresher registry instead of a second hard-coded list")
require.Contains(t, normalized, "credentials ? 'refresh_token'")
require.Contains(t, normalized, "btrim(credentials->>'refresh_token') <> ''")
require.Contains(t, normalized, "temp_unschedulable_until > NOW()")
@@ -152,14 +290,52 @@ func TestAccountRepository_ListOAuthRefreshCandidates_SQLFilter(t *testing.T) {
"must use IS NOT TRUE so accounts with NULL temp_unschedulable_until are not silently excluded by PG 3-valued logic")
require.NotContains(t, normalized, "AND NOT (",
"plain NOT (...) excludes NULL temp_unschedulable_until rows (the common healthy case)")
require.Contains(t, normalized, "ORDER BY priority ASC, id ASC")
require.Contains(t, normalized, "id > $2")
require.Contains(t, normalized, "ORDER BY id ASC")
require.Contains(t, normalized, "LIMIT $3")
require.NotContains(t, normalized, "credentials->>'expires_at'")
require.Len(t, capturedArgs, 3)
require.Equal(t, int64(100), capturedArgs[1])
require.Equal(t, 200, capturedArgs[2])
valuer, ok := capturedArgs[0].(interface{ Value() (driver.Value, error) })
require.True(t, ok)
platforms, err := valuer.Value()
require.NoError(t, err)
require.Contains(t, platforms, service.PlatformGrok)
require.NoError(t, mock.ExpectationsWereMet())
}
func TestAccountRepository_ListOAuthRefreshCandidatePage_ReconciliationExcludesAPIKeys(t *testing.T) {
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
require.NoError(t, err)
defer func() { _ = db.Close() }()
var capturedSQL string
mock.ExpectQuery("SELECT id").WillReturnRows(sqlmock.NewRows([]string{"id"}))
repo := newAccountRepositoryWithSQL(nil, captureQuerySQL{db: db, captured: &capturedSQL}, nil)
page, err := repo.ListOAuthRefreshCandidatePage(context.Background(), service.OAuthRefreshPageOptions{
Platforms: []string{service.PlatformGrok},
AfterID: 0,
Limit: 50,
})
require.NoError(t, err)
require.Empty(t, page.Accounts)
normalized := normalizeSQLWhitespace(capturedSQL)
require.Contains(t, normalized, "type = 'oauth'")
require.NotContains(t, normalized, "type IN ('oauth', 'setup-token')")
require.NotContains(t, normalized, "type = 'api-key'")
require.NotContains(t, normalized, "credentials ? 'refresh_token'",
"reconciliation must be able to find structurally invalid OAuth rows")
require.Contains(t, normalized, "ORDER BY id ASC")
require.NoError(t, mock.ExpectationsWereMet())
}
type captureQuerySQL struct {
db *sql.DB
captured *string
args *[]any
}
func (c captureQuerySQL) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
@@ -170,6 +346,9 @@ func (c captureQuerySQL) QueryContext(ctx context.Context, query string, args ..
if c.captured != nil {
*c.captured = query
}
if c.args != nil {
*c.args = append([]any(nil), args...)
}
return c.db.QueryContext(ctx, query, args...)
}
+1
View File
@@ -401,6 +401,7 @@ func registerGrokOAuthRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
grok.POST("/oauth/refresh-token", h.Admin.GrokOAuth.RefreshToken)
grok.POST("/oauth/create-from-oauth", h.Admin.GrokOAuth.CreateAccountFromOAuth)
grok.POST("/sso-to-oauth", h.Admin.GrokOAuth.CreateAccountsFromSSO)
grok.POST("/oauth/reconcile", h.Admin.GrokOAuth.ReconcileOAuthAccounts)
grok.POST("/accounts/:id/refresh", h.Admin.GrokOAuth.RefreshAccountToken)
grok.GET("/accounts/:id/quota", h.Admin.GrokOAuth.QueryQuota)
grok.POST("/accounts/:id/reset-quota", h.Admin.GrokOAuth.ResetQuota)
+29 -1
View File
@@ -18,6 +18,35 @@ var (
const AccountListGroupUngrouped int64 = -1
const AccountPrivacyModeUnsetFilter = "__unset__"
// OAuthRefreshPageOptions describes one bounded, cursor-stable scan of OAuth
// accounts. Candidate platforms are supplied by TokenRefreshService's refresher
// registry so repository eligibility cannot drift from registered providers.
type OAuthRefreshPageOptions struct {
Platforms []string
AfterID int64
Limit int
ActiveOnly bool
IncludeSetupToken bool
RequireRefreshToken bool
ExcludeRetryCooldown bool
}
// OAuthRefreshCandidatePage keeps cursor metadata from the raw SQL ID page.
// Hydration may legitimately lose a concurrently deleted row, but callers can
// still advance past the raw page without truncating or duplicating the scan.
type OAuthRefreshCandidatePage struct {
Accounts []Account
NextAfterID int64
HasMore bool
}
// OAuthRefreshCandidatePager is intentionally narrower than AccountRepository.
// Production refresh cycles fail closed when the repository does not implement
// this bounded contract instead of silently falling back to an unpaged scan.
type OAuthRefreshCandidatePager interface {
ListOAuthRefreshCandidatePage(ctx context.Context, options OAuthRefreshPageOptions) (*OAuthRefreshCandidatePage, error)
}
type AccountRepository interface {
Create(ctx context.Context, account *Account) error
GetByID(ctx context.Context, id int64) (*Account, error)
@@ -44,7 +73,6 @@ type AccountRepository interface {
ListAllWithFilters(ctx context.Context, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, error)
ListByGroup(ctx context.Context, groupID int64) ([]Account, error)
ListActive(ctx context.Context) ([]Account, error)
ListOAuthRefreshCandidates(ctx context.Context) ([]Account, error)
ListByPlatform(ctx context.Context, platform string) ([]Account, error)
UpdateLastUsed(ctx context.Context, id int64) error
@@ -247,6 +247,8 @@ func classifyGrokCredentialFailure(account *Account, err error) grokCredentialFa
}
return false
}
var providerConfigErr *providerConfigurationRefreshError
var containmentErr *providerCycleContainmentRefreshError
switch {
case errors.Is(err, errGrokOAuthRefreshTokenMissing), errors.Is(err, errGrokOAuthAccessTokenMissing), errors.Is(err, errGrokOAuthAccessTokenExpired):
@@ -261,8 +263,12 @@ func classifyGrokCredentialFailure(account *Account, err error) grokCredentialFa
return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderDown, action: NextAccountStop, message: "Grok OAuth account state is temporarily unavailable"}
case errors.Is(err, errOAuthRefreshCredentialPersist):
return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderDown, action: NextAccountStop, message: "Grok OAuth shared credential state is temporarily unavailable"}
case errors.As(err, &containmentErr):
return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderDown, action: NextAccountStop, message: "Grok OAuth provider state is temporarily unavailable"}
case errors.Is(err, errOAuthRefreshAccountStateChanged):
return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonAccountChanged, action: NextAccountRetry, message: "Grok OAuth account eligibility changed"}
case errors.As(err, &providerConfigErr):
return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderConfig, action: NextAccountStop, message: "Grok OAuth provider configuration is unavailable"}
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"):
return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderConfig, action: NextAccountStop, message: "Grok OAuth provider configuration is unavailable"}
case contains("grok_oauth_proxy_lookup_failed"),
@@ -1546,7 +1546,7 @@ func TestCredentialFailureCASMissDoesNotRecoverIneligibleLatestCredential(t *tes
func TestGetRequestCredentialSharedCredentialPersistenceFailureStopsWithoutAccountMutation(t *testing.T) {
account := expiredGrokOAuthAccountForCredentialTest(782)
repo := &tokenRefreshAccountRepo{updateErr: errors.New("database unavailable")}
repo := &tokenRefreshAccountRepo{conditionalSuccessErr: errors.New("database unavailable")}
repo.accountsByID = map[int64]*Account{account.ID: account}
cache := &grokTokenCacheForProviderTest{lockResult: true}
provider := NewGrokTokenProvider(repo, cache)
@@ -1565,7 +1565,8 @@ func TestGetRequestCredentialSharedCredentialPersistenceFailureStopsWithoutAccou
require.Equal(t, GatewayFailureScopeProvider, failoverErr.Scope)
require.Equal(t, GrokCredentialReasonProviderDown, failoverErr.Reason)
require.Equal(t, NextAccountStop, failoverErr.NextAccountAction)
require.Equal(t, 1, repo.updateCredentialsCalls)
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)
@@ -0,0 +1,345 @@
package service
import (
"context"
"errors"
"fmt"
"strings"
"time"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
)
const (
defaultGrokOAuthReconcilePageSize = 50
maxGrokOAuthReconcilePageSize = 500
maxGrokOAuthReconcileWindow = 24 * time.Hour
GrokOAuthReconcileReasonMissingRefreshToken = "missing_refresh_token"
GrokOAuthReconcileReasonMissingAccessToken = "missing_access_token"
GrokOAuthReconcileReasonMissingExpiry = "missing_expiry"
GrokOAuthReconcileReasonInvalidExpiry = "invalid_expiry"
GrokOAuthReconcileReasonNearExpiry = "near_expiry"
GrokOAuthReconcileReasonCredentialRejected = "credential_rejected"
GrokOAuthReconcileActionBlock = "block_account"
GrokOAuthReconcileActionRefresh = "refresh_credentials"
GrokOAuthReconcileOutcomePlanned = "planned"
GrokOAuthReconcileOutcomeApplied = "applied"
GrokOAuthReconcileOutcomeSkipped = "skipped"
GrokOAuthReconcileOutcomeFailed = "failed"
GrokOAuthReconcileOutcomePartial = "partial"
)
var (
ErrGrokOAuthReconcileMode = infraerrors.BadRequest(
"GROK_OAUTH_RECONCILE_MODE_INVALID",
"apply requires dry_run=false and apply=true",
)
ErrGrokOAuthReconcileCursor = infraerrors.BadRequest(
"GROK_OAUTH_RECONCILE_CURSOR_INVALID",
"after_id must be non-negative",
)
ErrGrokOAuthReconcileLimit = infraerrors.BadRequest(
"GROK_OAUTH_RECONCILE_LIMIT_INVALID",
"limit is outside the allowed reconciliation page range",
)
ErrGrokOAuthReconcileWindow = infraerrors.BadRequest(
"GROK_OAUTH_RECONCILE_WINDOW_INVALID",
"refresh_window_seconds is outside the allowed range",
)
)
// GrokOAuthReconciler is the narrow admin-facing reconciliation port.
type GrokOAuthReconciler interface {
ReconcileGrokOAuth(ctx context.Context, input GrokOAuthReconcileInput) (*GrokOAuthReconcileResult, error)
}
// GrokOAuthConditionalErrorRepository is the narrow compare-and-set mutation
// used by reconciliation. The repository must only transition an active Grok
// OAuth account when its credential document still exactly matches the state
// observed immediately before the mutation.
type GrokOAuthConditionalErrorRepository interface {
SetGrokOAuthErrorIfCredentialsUnchanged(ctx context.Context, id int64, expectedCredentials map[string]any, errorMsg string) (bool, error)
}
type GrokOAuthReconcileInput struct {
DryRun bool
Apply bool
AfterID int64
Limit int
RefreshWindow time.Duration
}
// GrokOAuthReconcileItem is deliberately metadata-only. Credentials, account
// identity fields, provider response bodies, and raw errors never cross this API.
type GrokOAuthReconcileItem struct {
AccountID int64 `json:"account_id"`
Reason string `json:"reason"`
Action string `json:"action"`
Outcome string `json:"outcome"`
}
type GrokOAuthReconcileResult struct {
DryRun bool `json:"dry_run"`
Scanned int `json:"scanned"`
Actionable int `json:"actionable"`
WouldBlock int `json:"would_block"`
WouldRefresh int `json:"would_refresh"`
Blocked int `json:"blocked"`
Refreshed int `json:"refreshed"`
Skipped int `json:"skipped"`
Failed int `json:"failed"`
Partial int `json:"partial"`
Items []GrokOAuthReconcileItem `json:"items"`
NextAfterID int64 `json:"next_after_id"`
HasMore bool `json:"has_more"`
}
func (s *TokenRefreshService) ReconcileGrokOAuth(ctx context.Context, input GrokOAuthReconcileInput) (*GrokOAuthReconcileResult, error) {
if ctx == nil {
ctx = context.Background()
}
if input.Apply && input.DryRun {
return nil, ErrGrokOAuthReconcileMode
}
if input.AfterID < 0 {
return nil, ErrGrokOAuthReconcileCursor
}
limit := input.Limit
maxPageSize := s.grokOAuthReconcileMaxPageSize()
if limit == 0 {
limit = min(defaultGrokOAuthReconcilePageSize, maxPageSize)
}
if limit < 1 || limit > maxPageSize {
return nil, ErrGrokOAuthReconcileLimit
}
refreshWindow := input.RefreshWindow
if refreshWindow == 0 {
refreshWindow = grokTokenRefreshSkew
}
if refreshWindow < 0 || refreshWindow > maxGrokOAuthReconcileWindow {
return nil, ErrGrokOAuthReconcileWindow
}
if refreshWindow < grokTokenRefreshSkew {
refreshWindow = grokTokenRefreshSkew
}
dryRun := !input.Apply
pager := s.candidatePager
if pager == nil {
pager, _ = s.accountRepo.(OAuthRefreshCandidatePager)
}
if pager == nil {
return nil, errors.New("OAuth refresh candidate pager is not configured")
}
page, err := pager.ListOAuthRefreshCandidatePage(ctx, OAuthRefreshPageOptions{
Platforms: []string{PlatformGrok},
AfterID: input.AfterID,
Limit: limit,
ActiveOnly: true,
// Reconciliation scans OAuth only and intentionally does not require a
// refresh token so structurally invalid rows remain discoverable.
IncludeSetupToken: false,
RequireRefreshToken: false,
})
if err != nil {
return nil, err
}
if page == nil {
return nil, errors.New("OAuth reconciliation repository returned a nil cursor page")
}
accounts := page.Accounts
if !isStrictlyIncreasingAccountPage(accounts, input.AfterID) {
return nil, errors.New("OAuth reconciliation repository returned an invalid cursor page")
}
result := &GrokOAuthReconcileResult{
DryRun: dryRun,
Scanned: len(accounts),
Items: make([]GrokOAuthReconcileItem, 0, len(accounts)),
HasMore: page.HasMore,
}
if result.HasMore {
if page.NextAfterID <= input.AfterID {
return nil, errors.New("OAuth reconciliation repository returned invalid cursor metadata")
}
result.NextAfterID = page.NextAfterID
}
registration, ok := s.grokRegistration()
if !ok {
return nil, errors.New("grok OAuth refresher is not registered")
}
conditionalErrorRepo, supportsConditionalError := s.accountRepo.(GrokOAuthConditionalErrorRepository)
if input.Apply && !supportsConditionalError {
return nil, errors.New("grok OAuth conditional error mutation is not configured")
}
providerState := &tokenRefreshProviderState{
service: s,
registration: registration,
rateGate: s.providerRateGate(PlatformGrok),
poolGate: s.providerConcurrencyGate(PlatformGrok),
}
for i := range accounts {
if err := ctx.Err(); err != nil {
return nil, err
}
account := &accounts[i]
reason, action, actionable := classifyGrokOAuthReconcileAccount(account, refreshWindow)
if !actionable {
result.Skipped++
continue
}
result.Actionable++
item := GrokOAuthReconcileItem{
AccountID: account.ID,
Reason: reason,
Action: action,
Outcome: GrokOAuthReconcileOutcomePlanned,
}
if action == GrokOAuthReconcileActionBlock {
result.WouldBlock++
} else {
result.WouldRefresh++
}
if dryRun {
result.Items = append(result.Items, item)
continue
}
switch action {
case GrokOAuthReconcileActionBlock:
latest, err := s.accountRepo.GetByID(ctx, account.ID)
if err != nil || latest == nil {
item.Outcome = GrokOAuthReconcileOutcomeFailed
result.Failed++
break
}
latestReason, latestAction, stillActionable := classifyGrokOAuthReconcileAccount(latest, refreshWindow)
if !stillActionable || latestAction != GrokOAuthReconcileActionBlock {
// The account changed after page hydration (for example, an admin
// reauthorized it). Never apply a stale destructive action; the next
// resumable scan can plan the fresh state.
item.Outcome = GrokOAuthReconcileOutcomeSkipped
result.Skipped++
break
}
account = latest
item.Reason = latestReason
applied, err := conditionalErrorRepo.SetGrokOAuthErrorIfCredentialsUnchanged(
ctx,
account.ID,
account.Credentials,
"Grok OAuth credential reconciliation: missing refresh token",
)
if err != nil {
item.Outcome = GrokOAuthReconcileOutcomeFailed
result.Failed++
break
}
if !applied {
// Reauthorization won the compare-and-set race after the final
// reread. The runtime fast path is installed only after the CAS
// succeeds, so the fresh active account remains untouched.
item.Outcome = GrokOAuthReconcileOutcomeSkipped
result.Skipped++
break
}
s.notifyAccountSchedulingBlocked(account, time.Time{}, "grok_oauth_reconcile_invalid")
account.Status = StatusError
account.Schedulable = false
cacheInvalidationFailed := s.cacheInvalidator == nil
if s.cacheInvalidator != nil {
if err := s.cacheInvalidator.InvalidateToken(ctx, account); err != nil {
cacheInvalidationFailed = true
}
}
result.Blocked++
if cacheInvalidationFailed {
item.Outcome = GrokOAuthReconcileOutcomePartial
result.Partial++
} else {
item.Outcome = GrokOAuthReconcileOutcomeApplied
}
case GrokOAuthReconcileActionRefresh:
if providerState.isTripped() {
item.Outcome = GrokOAuthReconcileOutcomeSkipped
result.Skipped++
break
}
if providerState.isTripped() {
item.Outcome = GrokOAuthReconcileOutcomeSkipped
result.Skipped++
break
}
refreshErr := s.refreshWithRetryWithRateGate(ctx, account, registration.refresher, registration.executor, refreshWindow, providerState)
providerState.recordResult(refreshErr)
var permanentErr *accountPermanentRefreshError
switch {
case refreshErr == nil:
item.Outcome = GrokOAuthReconcileOutcomeApplied
result.Refreshed++
case errors.Is(refreshErr, errRefreshSkipped):
item.Outcome = GrokOAuthReconcileOutcomeSkipped
result.Skipped++
case errors.As(refreshErr, &permanentErr) && permanentErr.persistentlyBlocked:
item.Reason = GrokOAuthReconcileReasonCredentialRejected
item.Action = GrokOAuthReconcileActionBlock
result.Blocked++
if permanentErr.cacheInvalidationFailed {
item.Outcome = GrokOAuthReconcileOutcomePartial
result.Partial++
} else {
item.Outcome = GrokOAuthReconcileOutcomeApplied
}
default:
item.Outcome = GrokOAuthReconcileOutcomeFailed
result.Failed++
}
default:
return nil, fmt.Errorf("unsupported Grok OAuth reconciliation action")
}
result.Items = append(result.Items, item)
}
return result, nil
}
func (s *TokenRefreshService) grokOAuthReconcileMaxPageSize() int {
return maxGrokOAuthReconcilePageSize
}
func (s *TokenRefreshService) grokRegistration() (tokenRefreshRegistration, bool) {
for _, registration := range s.registrations {
if registration.platform == PlatformGrok && registration.refresher != nil {
return registration, true
}
}
return tokenRefreshRegistration{}, false
}
func classifyGrokOAuthReconcileAccount(account *Account, refreshWindow time.Duration) (reason, action string, actionable bool) {
if account == nil || !account.IsGrokOAuth() || account.Status != StatusActive {
return "", "", false
}
if strings.TrimSpace(account.GetGrokRefreshToken()) == "" {
return GrokOAuthReconcileReasonMissingRefreshToken, GrokOAuthReconcileActionBlock, true
}
if strings.TrimSpace(account.GetGrokAccessToken()) == "" {
return GrokOAuthReconcileReasonMissingAccessToken, GrokOAuthReconcileActionRefresh, true
}
rawExpiry := strings.TrimSpace(account.GetCredential("expires_at"))
if rawExpiry == "" {
return GrokOAuthReconcileReasonMissingExpiry, GrokOAuthReconcileActionRefresh, true
}
expiresAt := account.GetCredentialAsTime("expires_at")
if expiresAt == nil {
return GrokOAuthReconcileReasonInvalidExpiry, GrokOAuthReconcileActionRefresh, true
}
if time.Until(*expiresAt) <= refreshWindow {
return GrokOAuthReconcileReasonNearExpiry, GrokOAuthReconcileActionRefresh, true
}
return "", "", false
}
@@ -0,0 +1,578 @@
package service
import (
"context"
"encoding/json"
"errors"
"reflect"
"sort"
"strings"
"sync"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/stretchr/testify/require"
)
type grokReconcileRepo struct {
AccountRepository
mu sync.Mutex
accounts []Account
requests []OAuthRefreshPageOptions
setErrorIDs []int64
updatedCredIDs []int64
setErrorMessage []string
getByIDOverrides map[int64]Account
pageOverride *OAuthRefreshCandidatePage
reauthorizeOnCAS bool
reauthorizeOnRefreshCAS bool
conditionalCalls int
}
func (r *grokReconcileRepo) GetByID(_ context.Context, id int64) (*Account, error) {
r.mu.Lock()
defer r.mu.Unlock()
if override, ok := r.getByIDOverrides[id]; ok {
account := override
return &account, nil
}
for i := range r.accounts {
if r.accounts[i].ID == id {
account := r.accounts[i]
return &account, nil
}
}
return nil, ErrAccountNotFound
}
func (r *grokReconcileRepo) ListOAuthRefreshCandidatePage(_ context.Context, options OAuthRefreshPageOptions) (*OAuthRefreshCandidatePage, error) {
r.mu.Lock()
defer r.mu.Unlock()
r.requests = append(r.requests, options)
if r.pageOverride != nil {
page := *r.pageOverride
page.Accounts = append([]Account(nil), r.pageOverride.Accounts...)
return &page, nil
}
accounts := append([]Account(nil), r.accounts...)
sort.Slice(accounts, func(i, j int) bool { return accounts[i].ID < accounts[j].ID })
page := make([]Account, 0, options.Limit)
for _, account := range accounts {
if account.ID <= options.AfterID {
continue
}
platformAllowed := false
for _, platform := range options.Platforms {
if account.Platform == platform {
platformAllowed = true
break
}
}
if !platformAllowed || options.ActiveOnly && account.Status != StatusActive {
continue
}
if options.IncludeSetupToken {
if account.Type != AccountTypeOAuth && account.Type != AccountTypeSetupToken {
continue
}
} else if account.Type != AccountTypeOAuth {
continue
}
if options.RequireRefreshToken && strings.TrimSpace(account.GetGrokRefreshToken()) == "" {
continue
}
page = append(page, account)
if len(page) == options.Limit {
break
}
}
result := &OAuthRefreshCandidatePage{Accounts: page, HasMore: len(page) == options.Limit}
if len(page) > 0 {
result.NextAfterID = page[len(page)-1].ID
}
return result, nil
}
func (r *grokReconcileRepo) UpdateCredentials(_ context.Context, id int64, credentials map[string]any) error {
r.mu.Lock()
defer r.mu.Unlock()
r.updatedCredIDs = append(r.updatedCredIDs, id)
for i := range r.accounts {
if r.accounts[i].ID == id {
r.accounts[i].Credentials = MergeCredentials(r.accounts[i].Credentials, credentials)
}
}
return nil
}
func (r *grokReconcileRepo) SetError(_ context.Context, id int64, message string) error {
r.mu.Lock()
defer r.mu.Unlock()
r.setErrorIDs = append(r.setErrorIDs, id)
r.setErrorMessage = append(r.setErrorMessage, message)
for i := range r.accounts {
if r.accounts[i].ID == id {
r.accounts[i].Status = StatusError
r.accounts[i].Schedulable = false
r.accounts[i].ErrorMessage = message
}
}
return nil
}
func (r *grokReconcileRepo) SetGrokOAuthErrorIfCredentialsUnchanged(
_ context.Context,
id int64,
expectedCredentials map[string]any,
message string,
) (bool, error) {
r.mu.Lock()
defer r.mu.Unlock()
r.conditionalCalls++
for i := range r.accounts {
account := &r.accounts[i]
if account.ID != id {
continue
}
if r.reauthorizeOnCAS {
r.reauthorizeOnCAS = false
account.Credentials = map[string]any{
"access_token": "fresh-access",
"refresh_token": "fresh-refresh",
"expires_at": time.Now().UTC().Add(4 * time.Hour).Format(time.RFC3339),
"_token_version": int64(2),
}
}
if account.Platform != PlatformGrok || account.Type != AccountTypeOAuth || account.Status != StatusActive ||
strings.TrimSpace(account.GetGrokRefreshToken()) != "" || !reflect.DeepEqual(account.Credentials, expectedCredentials) {
return false, nil
}
r.setErrorIDs = append(r.setErrorIDs, id)
r.setErrorMessage = append(r.setErrorMessage, message)
account.Status = StatusError
account.Schedulable = false
account.ErrorMessage = message
return true, nil
}
return false, nil
}
func (r *grokReconcileRepo) SetGrokOAuthRefreshErrorIfCredentialsUnchanged(
_ context.Context,
id int64,
expectedCredentials map[string]any,
expectedProxyID *int64,
message string,
) (bool, error) {
r.mu.Lock()
defer r.mu.Unlock()
for i := range r.accounts {
account := &r.accounts[i]
if account.ID != id {
continue
}
if r.reauthorizeOnRefreshCAS {
r.reauthorizeOnRefreshCAS = false
account.Credentials = map[string]any{
"access_token": "fresh-access",
"refresh_token": "fresh-refresh",
"expires_at": time.Now().UTC().Add(4 * time.Hour).Format(time.RFC3339),
"_token_version": int64(3),
}
}
if account.Platform != PlatformGrok || account.Type != AccountTypeOAuth || account.Status != StatusActive ||
!reflect.DeepEqual(account.ProxyID, expectedProxyID) ||
!reflect.DeepEqual(account.Credentials, expectedCredentials) {
return false, nil
}
r.setErrorIDs = append(r.setErrorIDs, id)
r.setErrorMessage = append(r.setErrorMessage, message)
account.Status = StatusError
account.Schedulable = false
account.ErrorMessage = message
return true, nil
}
return false, nil
}
func (r *grokReconcileRepo) SetGrokOAuthRefreshTempUnschedulableIfCredentialsUnchanged(
_ context.Context,
id int64,
expectedCredentials map[string]any,
expectedProxyID *int64,
until time.Time,
reason string,
) (bool, error) {
r.mu.Lock()
defer r.mu.Unlock()
for i := range r.accounts {
account := &r.accounts[i]
if account.ID != id {
continue
}
if account.Platform != PlatformGrok || account.Type != AccountTypeOAuth || account.Status != StatusActive ||
!reflect.DeepEqual(account.ProxyID, expectedProxyID) ||
!reflect.DeepEqual(account.Credentials, expectedCredentials) {
return false, nil
}
account.TempUnschedulableUntil = &until
account.TempUnschedulableReason = reason
return true, nil
}
return false, nil
}
func (r *grokReconcileRepo) snapshot() ([]OAuthRefreshPageOptions, []int64, []int64, []string) {
r.mu.Lock()
defer r.mu.Unlock()
return append([]OAuthRefreshPageOptions(nil), r.requests...), append([]int64(nil), r.setErrorIDs...), append([]int64(nil), r.updatedCredIDs...), append([]string(nil), r.setErrorMessage...)
}
type reconcileInvalidator struct {
mu sync.Mutex
ids []int64
err error
}
type reconcileRuntimeBlocker struct {
mu sync.Mutex
blocked []int64
cleared []int64
}
func (b *reconcileRuntimeBlocker) BlockAccountScheduling(account *Account, _ time.Time, _ string) {
b.mu.Lock()
defer b.mu.Unlock()
if account != nil {
b.blocked = append(b.blocked, account.ID)
}
}
func (b *reconcileRuntimeBlocker) ClearAccountSchedulingBlock(accountID int64) {
b.mu.Lock()
defer b.mu.Unlock()
b.cleared = append(b.cleared, accountID)
}
func (b *reconcileRuntimeBlocker) snapshot() (blocked, cleared []int64) {
b.mu.Lock()
defer b.mu.Unlock()
return append([]int64(nil), b.blocked...), append([]int64(nil), b.cleared...)
}
func (i *reconcileInvalidator) InvalidateToken(_ context.Context, account *Account) error {
i.mu.Lock()
defer i.mu.Unlock()
i.ids = append(i.ids, account.ID)
return i.err
}
func (i *reconcileInvalidator) count() int {
i.mu.Lock()
defer i.mu.Unlock()
return len(i.ids)
}
func newGrokReconcileService(repo *grokReconcileRepo, refresher *poolHealthRefresher, invalidator TokenCacheInvalidator) *TokenRefreshService {
return &TokenRefreshService{
accountRepo: repo,
candidatePager: repo,
cacheInvalidator: invalidator,
registrations: []tokenRefreshRegistration{{
platform: PlatformGrok,
refresher: refresher,
executor: refresher,
}},
refreshPolicy: DefaultBackgroundRefreshPolicy(),
cfg: &config.TokenRefreshConfig{
MaxRetries: 1,
CandidatePageSize: 50,
ProviderConcurrency: 2,
ProviderQPS: 100,
ProviderFailureThreshold: 3,
AttemptTimeoutSeconds: 1,
},
}
}
func grokReconcileFixtures() []Account {
now := time.Now().UTC()
return []Account{
{
ID: 1,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Credentials: map[string]any{"access_token": "access-secret"},
},
{
ID: 2,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Credentials: map[string]any{"refresh_token": "refresh-secret", "expires_at": now.Add(10 * time.Minute).Format(time.RFC3339)},
},
{
ID: 3,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Credentials: map[string]any{"access_token": "access-secret", "refresh_token": "refresh-secret", "expires_at": now.Add(30 * time.Minute).Format(time.RFC3339)},
},
{
ID: 4,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Credentials: map[string]any{"access_token": "access-secret", "refresh_token": "refresh-secret", "expires_at": now.Add(4 * time.Hour).Format(time.RFC3339)},
},
{
ID: 5,
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Status: StatusActive,
Schedulable: true,
Credentials: map[string]any{"api_key": "api-key-secret"},
},
}
}
func TestTokenRefreshService_ReconcileGrokOAuthDefaultsToDryRunAndSanitizedPlan(t *testing.T) {
repo := &grokReconcileRepo{accounts: grokReconcileFixtures()}
refresher := &poolHealthRefresher{}
svc := newGrokReconcileService(repo, refresher, nil)
result, err := svc.ReconcileGrokOAuth(context.Background(), GrokOAuthReconcileInput{})
require.NoError(t, err)
require.True(t, result.DryRun)
require.Equal(t, 4, result.Scanned, "Grok API-key rows must not enter the OAuth reconciliation page")
require.Equal(t, 3, result.Actionable)
require.Equal(t, 1, result.WouldBlock)
require.Equal(t, 2, result.WouldRefresh)
require.Zero(t, result.Blocked)
require.Zero(t, result.Refreshed)
require.Zero(t, refresher.calls.Load())
_, setErrorIDs, updatedIDs, _ := repo.snapshot()
require.Empty(t, setErrorIDs)
require.Empty(t, updatedIDs)
payload, err := json.Marshal(result)
require.NoError(t, err)
text := string(payload)
require.NotContains(t, text, "access-secret")
require.NotContains(t, text, "refresh-secret")
require.NotContains(t, text, "api-key-secret")
require.NotContains(t, text, `"credentials":`)
}
func TestGrokTokenRefresher_NeedsRefreshWhenAccessTokenMissingDespiteFarFutureExpiry(t *testing.T) {
refresher := NewGrokTokenRefresher(nil)
account := grokPoolAccount(99)
delete(account.Credentials, "access_token")
account.Credentials["expires_at"] = time.Now().UTC().Add(12 * time.Hour).Format(time.RFC3339)
require.True(t, refresher.NeedsRefresh(&account, time.Hour))
}
func TestTokenRefreshService_ReconcileGrokOAuthApplyIsIdempotent(t *testing.T) {
repo := &grokReconcileRepo{accounts: grokReconcileFixtures()}
invalidator := &reconcileInvalidator{}
refresher := &poolHealthRefresher{newCredentials: map[string]any{
"access_token": "rotated-access-secret",
"refresh_token": "rotated-refresh-secret",
"expires_at": time.Now().UTC().Add(4 * time.Hour).Format(time.RFC3339),
}}
svc := newGrokReconcileService(repo, refresher, invalidator)
first, err := svc.ReconcileGrokOAuth(context.Background(), GrokOAuthReconcileInput{Apply: true, Limit: 50})
require.NoError(t, err)
require.False(t, first.DryRun)
require.Equal(t, 1, first.Blocked)
require.Equal(t, 2, first.Refreshed)
require.Zero(t, first.Failed)
requests, setErrorIDs, updatedIDs, messages := repo.snapshot()
require.Equal(t, []int64{1}, setErrorIDs)
sort.Slice(updatedIDs, func(i, j int) bool { return updatedIDs[i] < updatedIDs[j] })
require.Equal(t, []int64{2, 3}, updatedIDs)
require.Len(t, messages, 1)
require.NotContains(t, messages[0], "secret")
require.False(t, requests[0].RequireRefreshToken, "structurally invalid rows must remain discoverable")
require.False(t, requests[0].IncludeSetupToken)
require.Equal(t, 3, invalidator.count(), "block and refresh actions must invalidate token cache state")
second, err := svc.ReconcileGrokOAuth(context.Background(), GrokOAuthReconcileInput{Apply: true, Limit: 50})
require.NoError(t, err)
require.Zero(t, second.Actionable)
require.Equal(t, int64(2), refresher.calls.Load(), "already refreshed rows must not be refreshed again")
_, setErrorIDs, updatedIDs, _ = repo.snapshot()
require.Equal(t, []int64{1}, setErrorIDs, "already blocked invalid rows must not transition twice")
require.Len(t, updatedIDs, 2)
}
func TestTokenRefreshService_ReconcileGrokOAuthCursorResumesWithoutDuplicates(t *testing.T) {
fixtures := grokReconcileFixtures()[:3]
repo := &grokReconcileRepo{accounts: fixtures}
svc := newGrokReconcileService(repo, &poolHealthRefresher{}, nil)
first, err := svc.ReconcileGrokOAuth(context.Background(), GrokOAuthReconcileInput{Limit: 2})
require.NoError(t, err)
require.True(t, first.HasMore)
require.Equal(t, int64(2), first.NextAfterID)
require.Len(t, first.Items, 2)
second, err := svc.ReconcileGrokOAuth(context.Background(), GrokOAuthReconcileInput{AfterID: first.NextAfterID, Limit: 2})
require.NoError(t, err)
require.False(t, second.HasMore)
require.Zero(t, second.NextAfterID)
require.Len(t, second.Items, 1)
require.NotEqual(t, first.Items[0].AccountID, second.Items[0].AccountID)
require.NotEqual(t, first.Items[1].AccountID, second.Items[0].AccountID)
}
func TestTokenRefreshService_ReconcileGrokOAuthCursorUsesRawPageAfterHydrationGap(t *testing.T) {
account := grokReconcileFixtures()[0]
repo := &grokReconcileRepo{pageOverride: &OAuthRefreshCandidatePage{
Accounts: []Account{account},
NextAfterID: account.ID + 1,
HasMore: true,
}}
svc := newGrokReconcileService(repo, &poolHealthRefresher{}, nil)
result, err := svc.ReconcileGrokOAuth(context.Background(), GrokOAuthReconcileInput{Limit: 2})
require.NoError(t, err)
require.True(t, result.HasMore)
require.Equal(t, account.ID+1, result.NextAfterID,
"cursor must advance past a raw selected ID that disappeared during hydration")
}
func TestTokenRefreshService_ReconcileGrokOAuthRejectsConflictingApplyMode(t *testing.T) {
svc := newGrokReconcileService(&grokReconcileRepo{}, &poolHealthRefresher{}, nil)
_, err := svc.ReconcileGrokOAuth(context.Background(), GrokOAuthReconcileInput{Apply: true, DryRun: true})
require.ErrorIs(t, err, ErrGrokOAuthReconcileMode)
}
func TestTokenRefreshService_ReconcileGrokOAuthSkipsStaleBlockAfterConcurrentReauthorization(t *testing.T) {
stale := grokReconcileFixtures()[0]
latest := stale
latest.Credentials = map[string]any{
"access_token": "fresh-access",
"refresh_token": "fresh-refresh",
"expires_at": time.Now().UTC().Add(4 * time.Hour).Format(time.RFC3339),
}
repo := &grokReconcileRepo{
accounts: []Account{stale},
getByIDOverrides: map[int64]Account{stale.ID: latest},
}
svc := newGrokReconcileService(repo, &poolHealthRefresher{}, &reconcileInvalidator{})
result, err := svc.ReconcileGrokOAuth(context.Background(), GrokOAuthReconcileInput{Apply: true, Limit: 50})
require.NoError(t, err)
require.Zero(t, result.Blocked)
require.Equal(t, 1, result.Skipped)
require.Equal(t, GrokOAuthReconcileOutcomeSkipped, result.Items[0].Outcome)
_, setErrorIDs, _, _ := repo.snapshot()
require.Empty(t, setErrorIDs, "a concurrently reauthorized account must not be disabled from stale page state")
}
func TestTokenRefreshService_ReconcileGrokOAuthDoesNotRuntimeBlockWhenReauthorizationWinsConditionalMutation(t *testing.T) {
account := grokReconcileFixtures()[0]
account.Credentials["_token_version"] = int64(1)
repo := &grokReconcileRepo{
accounts: []Account{account},
reauthorizeOnCAS: true,
}
invalidator := &reconcileInvalidator{}
blocker := &reconcileRuntimeBlocker{}
svc := newGrokReconcileService(repo, &poolHealthRefresher{}, invalidator)
svc.SetAccountRuntimeBlocker(blocker)
result, err := svc.ReconcileGrokOAuth(context.Background(), GrokOAuthReconcileInput{Apply: true, Limit: 50})
require.NoError(t, err)
require.Zero(t, result.Blocked)
require.Equal(t, 1, result.Skipped)
require.Equal(t, GrokOAuthReconcileOutcomeSkipped, result.Items[0].Outcome)
require.Zero(t, invalidator.count(), "a lost compare-and-set race must not invalidate fresh credentials")
_, setErrorIDs, _, _ := repo.snapshot()
require.Empty(t, setErrorIDs)
require.Equal(t, 1, repo.conditionalCalls)
blocked, cleared := blocker.snapshot()
require.Empty(t, blocked, "a lost compare-and-set race must never install a runtime block")
require.Empty(t, cleared, "reconciliation must not clear a block it does not own")
latest, getErr := repo.GetByID(context.Background(), account.ID)
require.NoError(t, getErr)
require.Equal(t, StatusActive, latest.Status)
require.True(t, latest.Schedulable)
require.Equal(t, "fresh-refresh", latest.GetGrokRefreshToken())
}
func TestTokenRefreshService_ReconcileGrokOAuthReportsPermanentRefreshMutationAsBlocked(t *testing.T) {
account := grokReconcileFixtures()[2]
repo := &grokReconcileRepo{accounts: []Account{account}}
refresher := &poolHealthRefresher{err: errors.New(`GROK_OAUTH_ENTITLEMENT_DENIED: subscription required`)}
svc := newGrokReconcileService(repo, refresher, &reconcileInvalidator{})
result, err := svc.ReconcileGrokOAuth(context.Background(), GrokOAuthReconcileInput{Apply: true, Limit: 50})
require.NoError(t, err)
require.Equal(t, 1, result.Blocked)
require.Zero(t, result.Failed)
require.Zero(t, result.Partial)
require.Equal(t, GrokOAuthReconcileActionBlock, result.Items[0].Action)
require.Equal(t, GrokOAuthReconcileReasonCredentialRejected, result.Items[0].Reason)
require.Equal(t, GrokOAuthReconcileOutcomeApplied, result.Items[0].Outcome)
_, setErrorIDs, _, _ := repo.snapshot()
require.Equal(t, []int64{account.ID}, setErrorIDs)
}
func TestTokenRefreshService_ReconcileGrokOAuthReportsConcurrentRefreshReauthorizationAsSkipped(t *testing.T) {
account := grokReconcileFixtures()[2]
repo := &grokReconcileRepo{
accounts: []Account{account},
reauthorizeOnRefreshCAS: true,
}
invalidator := &reconcileInvalidator{}
refresher := &poolHealthRefresher{err: errors.New("invalid_grant: revoked")}
svc := newGrokReconcileService(repo, refresher, invalidator)
result, err := svc.ReconcileGrokOAuth(context.Background(), GrokOAuthReconcileInput{Apply: true, Limit: 50})
require.NoError(t, err)
require.Equal(t, 1, result.Skipped)
require.Zero(t, result.Failed)
require.Zero(t, result.Blocked)
require.Equal(t, GrokOAuthReconcileOutcomeSkipped, result.Items[0].Outcome)
require.Zero(t, invalidator.count())
_, setErrorIDs, _, _ := repo.snapshot()
require.Empty(t, setErrorIDs)
latest, getErr := repo.GetByID(context.Background(), account.ID)
require.NoError(t, getErr)
require.Equal(t, StatusActive, latest.Status)
require.Equal(t, "fresh-refresh", latest.GetGrokRefreshToken())
}
func TestTokenRefreshService_ReconcileGrokOAuthReportsInvalidationFailureAsPartial(t *testing.T) {
account := grokReconcileFixtures()[0]
repo := &grokReconcileRepo{accounts: []Account{account}}
svc := newGrokReconcileService(repo, &poolHealthRefresher{}, &reconcileInvalidator{err: errors.New("cache unavailable")})
result, err := svc.ReconcileGrokOAuth(context.Background(), GrokOAuthReconcileInput{Apply: true, Limit: 50})
require.NoError(t, err)
require.Equal(t, 1, result.Blocked)
require.Equal(t, 1, result.Partial)
require.Zero(t, result.Failed)
require.Equal(t, GrokOAuthReconcileOutcomePartial, result.Items[0].Outcome)
}
@@ -30,6 +30,9 @@ func (r *GrokTokenRefresher) NeedsRefresh(account *Account, refreshWindow time.D
if account == nil || strings.TrimSpace(account.GetGrokRefreshToken()) == "" {
return false
}
if strings.TrimSpace(account.GetGrokAccessToken()) == "" {
return true
}
expiresAt := account.GetCredentialAsTime("expires_at")
if expiresAt == nil {
return true
+202 -47
View File
@@ -20,8 +20,25 @@ type OAuthRefreshExecutor interface {
CacheKey(account *Account) string
}
const defaultRefreshLockTTL = 60 * time.Second
const oauthRefreshLockCleanupTimeout = 2 * time.Second
// GrokOAuthRefreshSuccessRepository is the persistence boundary for a
// provider-issued Grok credential rotation. Implementations must compare the
// complete credential document and proxy used by the upstream attempt, and
// atomically publish scheduler invalidation with a successful update.
type GrokOAuthRefreshSuccessRepository interface {
UpdateGrokOAuthCredentialsIfUnchanged(
ctx context.Context,
id int64,
expectedCredentials map[string]any,
expectedProxyID *int64,
credentials map[string]any,
) (bool, error)
}
const (
defaultRefreshLockTTL = 60 * time.Second
defaultRefreshLockReleaseTimeout = 2 * time.Second
defaultRefreshPostPersistCleanupTimeout = 2 * time.Second
)
var (
errOAuthRefreshAccountRereadFailed = errors.New("oauth refresh account reread failed")
@@ -40,42 +57,78 @@ func isOAuthRefreshRequestPath(ctx context.Context) bool {
return requestPath
}
type oauthRefreshLocalLock struct {
semaphore chan struct{}
type contextMutex struct {
token chan struct{}
}
// Keep the request-path credential mutation lock API introduced by #4212
// while sharing the context-aware mutex implementation used by pool refresh.
type oauthRefreshLocalLock = contextMutex
func newOAuthRefreshLocalLock() *oauthRefreshLocalLock {
return &oauthRefreshLocalLock{semaphore: make(chan struct{}, 1)}
return newContextMutex()
}
func (l *oauthRefreshLocalLock) Lock(ctx context.Context) error {
type oauthRefreshStateUnavailableError struct {
err error
}
func (e *oauthRefreshStateUnavailableError) Error() string {
return "OAuth refresh account state is unavailable"
}
func (e *oauthRefreshStateUnavailableError) Unwrap() error {
if e == nil {
return nil
}
return e.err
}
func newContextMutex() *contextMutex {
return &contextMutex{token: make(chan struct{}, 1)}
}
func (m *contextMutex) Lock(ctx context.Context) error {
select {
case l.semaphore <- struct{}{}:
case m.token <- struct{}{}:
return nil
case <-ctx.Done():
return ctx.Err()
}
}
func (l *oauthRefreshLocalLock) Unlock() {
<-l.semaphore
func (m *contextMutex) Unlock() {
<-m.token
}
// OAuthRefreshResult 统一刷新结果
type OAuthRefreshResult struct {
Refreshed bool // 实际执行了刷新
NewCredentials map[string]any // 刷新后的 credentialsnil 表示未刷新)
Account *Account // 从 DB 重新读取的最新 account
Account *Account // 成功时为最新 account;刷新错误时为实际尝试的凭据快照
LockHeld bool // 锁被其他 worker 持有(未执行刷新)
}
func snapshotOAuthRefreshAccount(account *Account) *Account {
if account == nil {
return nil
}
snapshot := *account
snapshot.Credentials = shallowCopyMap(account.Credentials)
if account.ProxyID != nil {
proxyID := *account.ProxyID
snapshot.ProxyID = &proxyID
}
return &snapshot
}
// OAuthRefreshAPI 统一的 OAuth Token 刷新入口
// 封装分布式锁、进程内互斥锁、DB 重读、已刷新检查、竞争恢复等通用逻辑
type OAuthRefreshAPI struct {
accountRepo AccountRepository
tokenCache GeminiTokenCache // 可选,nil = 无分布式锁
lockTTL time.Duration
localLocks sync.Map // key: cacheKey string -> value: *oauthRefreshLocalLock
localLocks sync.Map // key: cacheKey string -> value: *contextMutex
}
// NewOAuthRefreshAPI 创建统一刷新 API
@@ -93,11 +146,11 @@ func NewOAuthRefreshAPI(accountRepo AccountRepository, tokenCache GeminiTokenCac
}
// getLocalLock 返回指定 cacheKey 的进程内互斥锁
func (api *OAuthRefreshAPI) getLocalLock(cacheKey string) *oauthRefreshLocalLock {
actual, _ := api.localLocks.LoadOrStore(cacheKey, newOAuthRefreshLocalLock())
mu, ok := actual.(*oauthRefreshLocalLock)
func (api *OAuthRefreshAPI) getLocalLock(cacheKey string) *contextMutex {
actual, _ := api.localLocks.LoadOrStore(cacheKey, newContextMutex())
mu, ok := actual.(*contextMutex)
if !ok {
mu = newOAuthRefreshLocalLock()
mu = newContextMutex()
api.localLocks.Store(cacheKey, mu)
}
return mu
@@ -127,6 +180,7 @@ func (api *OAuthRefreshAPI) RefreshIfNeeded(
if executor == nil {
return nil, errors.New("oauth refresh executor is nil")
}
requestPath := isOAuthRefreshRequestPath(ctx)
cacheKey := executor.CacheKey(account)
// 0. 获取进程内互斥锁(防止同一进程内的并发刷新竞争)
@@ -150,38 +204,46 @@ func (api *OAuthRefreshAPI) RefreshIfNeeded(
// 锁被其他 worker 持有
return &OAuthRefreshResult{LockHeld: true}, nil
} else {
defer func() {
cleanupCtx, cancel := context.WithTimeout(context.Background(), oauthRefreshLockCleanupTimeout)
defer cancel()
_ = api.tokenCache.ReleaseRefreshLock(cleanupCtx, cacheKey)
}()
defer api.releaseRefreshLock(ctx, cacheKey)
}
}
// 2. 从 DB 重读最新 account(锁保护下,确保使用最新的 refresh_token
freshAccount, err := api.accountRepo.GetByID(ctx, account.ID)
if err != nil {
return nil, fmt.Errorf("%w: %v", errOAuthRefreshAccountRereadFailed, err)
if requestPath {
return nil, fmt.Errorf("%w: %v", errOAuthRefreshAccountRereadFailed, err)
}
return nil, &oauthRefreshStateUnavailableError{err: err}
}
if freshAccount == nil {
return nil, fmt.Errorf("%w: account not found", errOAuthRefreshAccountStateChanged)
if requestPath {
return nil, fmt.Errorf("%w: account not found", errOAuthRefreshAccountStateChanged)
}
return nil, &oauthRefreshStateUnavailableError{err: fmt.Errorf("account not found")}
}
if freshAccount.ID != account.ID {
return nil, fmt.Errorf("%w: account identity mismatch", errOAuthRefreshAccountRereadFailed)
}
if !freshAccount.IsActive() {
return nil, fmt.Errorf("%w: account is not active", errOAuthRefreshAccountStateChanged)
if requestPath {
return nil, fmt.Errorf("%w: account is not active", errOAuthRefreshAccountStateChanged)
}
return &OAuthRefreshResult{Account: freshAccount}, nil
}
if isOAuthRefreshRequestPath(ctx) && freshAccount.Platform == PlatformGrok {
if requestPath && freshAccount.Platform == PlatformGrok {
if eligibilityErr := grokOAuthRequestAccountEligibilityError(freshAccount); eligibilityErr != nil {
return nil, withGrokCredentialFailureSnapshot(eligibilityErr, freshAccount)
}
}
if !executor.CanRefresh(freshAccount) {
if freshAccount.IsGrokOAuth() && strings.TrimSpace(freshAccount.GetGrokRefreshToken()) == "" {
if requestPath && freshAccount.IsGrokOAuth() && strings.TrimSpace(freshAccount.GetGrokRefreshToken()) == "" {
return nil, withGrokCredentialFailureSnapshot(errGrokOAuthRefreshTokenMissing, freshAccount)
}
return nil, fmt.Errorf("%w: account is no longer refreshable", errOAuthRefreshAccountStateChanged)
if requestPath {
return nil, fmt.Errorf("%w: account is no longer refreshable", errOAuthRefreshAccountStateChanged)
}
return &OAuthRefreshResult{Account: freshAccount}, nil
}
// 3. 二次检查是否仍需刷新(另一条路径可能已刷新)
@@ -192,16 +254,19 @@ func (api *OAuthRefreshAPI) RefreshIfNeeded(
}
// 4. 执行平台特定刷新逻辑
attemptedAccount := snapshotOAuthRefreshAccount(freshAccount)
newCredentials, refreshErr := executor.Refresh(ctx, freshAccount)
if err := ctx.Err(); err != nil {
return nil, err
if ctxErr := ctx.Err(); ctxErr != nil {
// A provider implementation may ignore cancellation and return late
// credentials. Never persist them after the attempt/cycle boundary.
return nil, ctxErr
}
if refreshErr != nil {
// 竞争恢复:invalid_grant 可能是另一个 worker 已消费了旧 refresh_token
// 重新读取 DB,如果 refresh_token 已更新则说明是竞争,返回成功
if isInvalidGrantError(refreshErr) {
if recoveredAccount, recovered := api.tryRecoverFromRefreshRace(ctx, freshAccount); recovered {
if isOAuthRefreshRequestPath(ctx) && recoveredAccount.Platform == PlatformGrok {
if requestPath && recoveredAccount.Platform == PlatformGrok {
if eligibilityErr := grokOAuthRequestAccountEligibilityError(recoveredAccount); eligibilityErr != nil {
return nil, withGrokCredentialFailureSnapshot(eligibilityErr, recoveredAccount)
}
@@ -215,44 +280,134 @@ func (api *OAuthRefreshAPI) RefreshIfNeeded(
}, nil
}
}
return nil, withGrokCredentialFailureSnapshot(refreshErr, freshAccount)
// Preserve the exact account snapshot used by the failed upstream call.
// Callers can then conditionally mutate only that credential version and
// avoid quarantining a concurrently reauthorized account.
result := &OAuthRefreshResult{Account: attemptedAccount}
if requestPath && attemptedAccount.Platform == PlatformGrok {
return result, withGrokCredentialFailureSnapshot(refreshErr, attemptedAccount)
}
return result, refreshErr
}
// 5. 设置版本号 + 更新 DB
if newCredentials != nil {
newCredentials["_token_version"] = time.Now().UnixMilli()
if updateErr := persistAccountCredentials(ctx, api.accountRepo, freshAccount, newCredentials); updateErr != nil {
if freshAccount.IsGrokOAuth() {
conditionalRepo, ok := api.accountRepo.(GrokOAuthRefreshSuccessRepository)
if !ok {
return nil, &providerConfigurationRefreshError{
err: fmt.Errorf("grok OAuth refresh success CAS repository is not configured"),
}
}
applied, updateErr := conditionalRepo.UpdateGrokOAuthCredentialsIfUnchanged(
ctx,
freshAccount.ID,
attemptedAccount.Credentials,
attemptedAccount.ProxyID,
newCredentials,
)
if updateErr != nil {
slog.Error("oauth_refresh_update_failed",
"account_id", freshAccount.ID,
"platform", freshAccount.Platform,
"error", updateErr,
)
// The provider may have rotated and consumed the refresh token.
// Retrying after an ambiguous local persistence result can turn a
// healthy account into invalid_grant, so contain this provider cycle.
return nil, &providerCycleContainmentRefreshError{
err: fmt.Errorf("OAuth refresh succeeded but credential persistence failed: %w", updateErr),
}
}
if !applied {
currentAccount, readErr := api.accountRepo.GetByID(ctx, freshAccount.ID)
if readErr != nil || currentAccount == nil {
if readErr == nil {
readErr = fmt.Errorf("account not found after Grok OAuth success CAS miss")
}
return nil, &providerCycleContainmentRefreshError{
err: fmt.Errorf("grok OAuth success CAS lost and current state is unavailable: %w", readErr),
}
}
slog.Info("oauth_refresh_success_cas_skipped_stale_credentials",
"account_id", freshAccount.ID,
"platform", freshAccount.Platform,
)
return &OAuthRefreshResult{Account: currentAccount}, nil
}
durableAccount, readErr := api.loadGrokDurableAccountAfterPersist(ctx, cacheKey, freshAccount.ID)
if readErr != nil || durableAccount == nil {
if readErr == nil {
readErr = fmt.Errorf("account not found after Grok OAuth success CAS")
}
return nil, &providerCycleContainmentRefreshError{
err: fmt.Errorf("grok OAuth success persisted but durable account state is unavailable: %w", readErr),
}
}
// The CAS changes credentials only. A concurrent admin or scheduler
// mutation may have changed status, schedulability, or cooldown fields
// while the provider call was in flight. Return the durable row so
// post-refresh cache publication cannot restore that stale snapshot.
freshAccount = durableAccount
} else if updateErr := persistAccountCredentials(ctx, api.accountRepo, freshAccount, newCredentials); updateErr != nil {
slog.Error("oauth_refresh_update_failed",
"account_id", freshAccount.ID,
"error", updateErr,
)
return nil, withGrokCredentialFailureSnapshot(
fmt.Errorf("%w: %v", errOAuthRefreshCredentialPersist, updateErr), freshAccount,
)
return nil, fmt.Errorf("%w: %v", errOAuthRefreshCredentialPersist, updateErr)
}
}
resultAccount := freshAccount
if isOAuthRefreshRequestPath(ctx) && freshAccount.Platform == PlatformGrok {
latestAccount, rereadErr := api.accountRepo.GetByID(ctx, freshAccount.ID)
if rereadErr != nil {
return nil, fmt.Errorf("%w: %v", errOAuthRefreshAccountRereadFailed, rereadErr)
if requestPath && freshAccount.Platform == PlatformGrok {
if eligibilityErr := grokOAuthRequestAccountEligibilityError(freshAccount); eligibilityErr != nil {
return nil, withGrokCredentialFailureSnapshot(eligibilityErr, freshAccount)
}
if latestAccount == nil {
return nil, fmt.Errorf("%w: account not found after refresh", errOAuthRefreshAccountStateChanged)
}
if eligibilityErr := grokOAuthRequestAccountEligibilityError(latestAccount); eligibilityErr != nil {
return nil, withGrokCredentialFailureSnapshot(eligibilityErr, latestAccount)
}
resultAccount = latestAccount
}
return &OAuthRefreshResult{
Refreshed: true,
NewCredentials: newCredentials,
Account: resultAccount,
Account: freshAccount,
}, nil
}
func (api *OAuthRefreshAPI) releaseRefreshLock(parent context.Context, cacheKey string) {
cleanupParent := context.Background()
if parent != nil {
cleanupParent = context.WithoutCancel(parent)
}
ctx, cancel := context.WithTimeout(cleanupParent, defaultRefreshLockReleaseTimeout)
defer cancel()
if err := api.tokenCache.ReleaseRefreshLock(ctx, cacheKey); err != nil {
slog.Warn("oauth_refresh_lock_release_failed", "cache_key", cacheKey, "error", err)
}
}
func (api *OAuthRefreshAPI) loadGrokDurableAccountAfterPersist(parent context.Context, cacheKey string, accountID int64) (*Account, error) {
cleanupParent := context.Background()
if parent != nil {
cleanupParent = context.WithoutCancel(parent)
}
ctx, cancel := context.WithTimeout(cleanupParent, defaultRefreshPostPersistCleanupTimeout)
defer cancel()
// A successful rotation can revoke the access token still cached from the
// pre-rotation credential document. Trigger deletion at the commit boundary,
// even if the attempt/parent context was canceled immediately after CAS.
if api.tokenCache != nil {
if err := api.tokenCache.DeleteAccessToken(ctx, cacheKey); err != nil {
slog.Warn("oauth_refresh_post_persist_cache_delete_failed",
"account_id", accountID,
"cache_key", cacheKey,
"error", err,
)
}
}
return api.accountRepo.GetByID(ctx, accountID)
}
// isInvalidGrantError 检查错误是否为 invalid_grant
func isInvalidGrantError(err error) bool {
return err != nil && strings.Contains(strings.ToLower(err.Error()), "invalid_grant")
@@ -5,6 +5,7 @@ package service
import (
"context"
"errors"
"reflect"
"sync"
"testing"
"time"
@@ -17,18 +18,38 @@ import (
// refreshAPIAccountRepo implements AccountRepository for OAuthRefreshAPI tests.
type refreshAPIAccountRepo struct {
mockAccountRepoForGemini
account *Account // returned by GetByID
getByIDErr error
updateErr error
updateCalls int
updateCredentialsCalls int
account *Account // returned by GetByID
getByIDErr error
getByIDCalls int
getByIDErrAfterCall int
getByIDErrAfterCallErr error
updateErr error
updateCalls int
updateCredentialsCalls int
successCASCalls int
beforeSuccessCAS func(*refreshAPIAccountRepo)
lastExpectedCredentials map[string]any
lastExpectedProxyID *int64
}
func (r *refreshAPIAccountRepo) GetByID(_ context.Context, _ int64) (*Account, error) {
r.getByIDCalls++
if r.getByIDErrAfterCall > 0 && r.getByIDCalls >= r.getByIDErrAfterCall {
return nil, r.getByIDErrAfterCallErr
}
if r.getByIDErr != nil {
return nil, r.getByIDErr
}
return r.account, nil
return activeRefreshAPITestAccount(r.account), nil
}
func activeRefreshAPITestAccount(account *Account) *Account {
if account == nil || account.Status != "" {
return account
}
copy := *account
copy.Status = StatusActive
return &copy
}
func (r *refreshAPIAccountRepo) Update(_ context.Context, _ *Account) error {
@@ -49,6 +70,39 @@ func (r *refreshAPIAccountRepo) UpdateCredentials(_ context.Context, id int64, c
return nil
}
func (r *refreshAPIAccountRepo) UpdateGrokOAuthCredentialsIfUnchanged(
_ context.Context,
id int64,
expectedCredentials map[string]any,
expectedProxyID *int64,
credentials map[string]any,
) (bool, error) {
r.successCASCalls++
r.lastExpectedCredentials = shallowCopyMap(expectedCredentials)
if expectedProxyID != nil {
proxyID := *expectedProxyID
r.lastExpectedProxyID = &proxyID
} else {
r.lastExpectedProxyID = nil
}
if r.beforeSuccessCAS != nil {
r.beforeSuccessCAS(r)
}
if r.updateErr != nil {
return false, r.updateErr
}
if r.account == nil || r.account.ID != id || r.account.Platform != PlatformGrok ||
r.account.Type != AccountTypeOAuth ||
!reflect.DeepEqual(r.account.Credentials, expectedCredentials) ||
!reflect.DeepEqual(r.account.ProxyID, expectedProxyID) {
return false, nil
}
r.updateCalls++
r.updateCredentialsCalls++
r.account.Credentials = shallowCopyMap(credentials)
return true, nil
}
// refreshAPIExecutorStub implements OAuthRefreshExecutor for tests.
type refreshAPIExecutorStub struct {
needsRefresh bool
@@ -56,9 +110,20 @@ type refreshAPIExecutorStub struct {
credentials map[string]any
err error
refreshCalls int
canRefresh func(*Account) bool
onRefresh func()
delay time.Duration
}
func (e *refreshAPIExecutorStub) CanRefresh(_ *Account) bool { return !e.cannotRefresh }
func (e *refreshAPIExecutorStub) CanRefresh(account *Account) bool {
if e.cannotRefresh {
return false
}
if e.canRefresh != nil {
return e.canRefresh(account)
}
return true
}
func (e *refreshAPIExecutorStub) NeedsRefresh(_ *Account, _ time.Duration) bool {
return e.needsRefresh
@@ -66,6 +131,12 @@ func (e *refreshAPIExecutorStub) NeedsRefresh(_ *Account, _ time.Duration) bool
func (e *refreshAPIExecutorStub) Refresh(_ context.Context, _ *Account) (map[string]any, error) {
e.refreshCalls++
if e.delay > 0 {
time.Sleep(e.delay)
}
if e.onRefresh != nil {
e.onRefresh()
}
if e.err != nil {
return nil, e.err
}
@@ -82,6 +153,9 @@ type refreshAPICacheStub struct {
lockErr error
releaseCalls int
releaseCtxErr error
deleteCalls int
deleteKey string
deleteCtxErr error
}
func (c *refreshAPICacheStub) GetAccessToken(context.Context, string) (string, error) {
@@ -92,7 +166,12 @@ func (c *refreshAPICacheStub) SetAccessToken(context.Context, string, string, ti
return nil
}
func (c *refreshAPICacheStub) DeleteAccessToken(context.Context, string) error { return nil }
func (c *refreshAPICacheStub) DeleteAccessToken(ctx context.Context, key string) error {
c.deleteCalls++
c.deleteKey = key
c.deleteCtxErr = ctx.Err()
return nil
}
func (c *refreshAPICacheStub) AcquireRefreshLock(context.Context, string, time.Duration) (bool, error) {
return c.lockResult, c.lockErr
@@ -236,7 +315,9 @@ func TestRefreshIfNeeded_RefreshError(t *testing.T) {
result, err := api.RefreshIfNeeded(context.Background(), account, executor, 3*time.Minute)
require.Error(t, err)
require.Nil(t, result)
require.NotNil(t, result)
require.NotNil(t, result.Account)
require.Equal(t, account.ID, result.Account.ID)
require.Contains(t, err.Error(), "invalid_grant")
require.Equal(t, 0, repo.updateCalls) // no DB update on refresh error
require.Equal(t, 1, cache.releaseCalls) // lock still released via defer
@@ -263,6 +344,123 @@ func TestRefreshIfNeeded_DBUpdateError(t *testing.T) {
require.Equal(t, 1, repo.updateCalls) // attempted
}
func TestRefreshIfNeeded_GrokSuccessCASLetsConcurrentReauthorizationWin(t *testing.T) {
proxyID := int64(17)
account := &Account{
ID: 70,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
ProxyID: &proxyID,
Credentials: map[string]any{
"access_token": "attempted-access",
"refresh_token": "attempted-refresh",
"_token_version": int64(1),
},
}
repo := &refreshAPIAccountRepo{account: account}
repo.beforeSuccessCAS = func(r *refreshAPIAccountRepo) {
repairedProxyID := int64(23)
r.account.ProxyID = &repairedProxyID
r.account.Credentials = map[string]any{
"access_token": "reauthorized-access",
"refresh_token": "reauthorized-refresh",
"_token_version": int64(2),
}
}
executor := &refreshAPIExecutorStub{
needsRefresh: true,
credentials: map[string]any{
"access_token": "provider-access",
"refresh_token": "provider-refresh",
},
}
result, err := NewOAuthRefreshAPI(repo, nil).RefreshIfNeeded(context.Background(), account, executor, time.Hour)
require.NoError(t, err)
require.NotNil(t, result)
require.False(t, result.Refreshed, "a lost success CAS is an already-refreshed skip")
require.Nil(t, result.NewCredentials)
require.Equal(t, "reauthorized-refresh", result.Account.GetGrokRefreshToken())
require.NotNil(t, result.Account.ProxyID)
require.Equal(t, int64(23), *result.Account.ProxyID)
require.Equal(t, 1, repo.successCASCalls)
require.Equal(t, "attempted-refresh", repo.lastExpectedCredentials["refresh_token"])
require.NotNil(t, repo.lastExpectedProxyID)
require.Equal(t, proxyID, *repo.lastExpectedProxyID)
require.Zero(t, repo.updateCredentialsCalls, "the provider result must not overwrite a concurrent repair")
}
func TestRefreshIfNeeded_GrokSuccessPersistenceFailureIsProviderContainment(t *testing.T) {
account := &Account{
ID: 71,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Credentials: map[string]any{
"access_token": "attempted-access",
"refresh_token": "attempted-refresh",
},
}
repo := &refreshAPIAccountRepo{account: account, updateErr: errors.New("database unavailable")}
executor := &refreshAPIExecutorStub{
needsRefresh: true,
credentials: map[string]any{
"access_token": "provider-access",
"refresh_token": "provider-refresh",
},
}
result, err := NewOAuthRefreshAPI(repo, nil).RefreshIfNeeded(context.Background(), account, executor, time.Hour)
require.Error(t, err)
require.Nil(t, result)
var containmentErr *providerCycleContainmentRefreshError
require.ErrorAs(t, err, &containmentErr)
require.Equal(t, "attempted-refresh", account.GetGrokRefreshToken(),
"an ambiguous persistence result must not mutate the in-memory account")
require.Equal(t, 1, repo.successCASCalls)
require.Zero(t, repo.updateCredentialsCalls)
}
func TestRefreshIfNeeded_GrokSuccessDurableRereadFailureIsProviderContainment(t *testing.T) {
account := &Account{
ID: 72,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Credentials: map[string]any{
"access_token": "attempted-access",
"refresh_token": "attempted-refresh",
},
}
repo := &refreshAPIAccountRepo{
account: account,
getByIDErrAfterCall: 2,
getByIDErrAfterCallErr: errors.New("durable state unavailable"),
}
cache := &refreshAPICacheStub{lockResult: true}
executor := &refreshAPIExecutorStub{
needsRefresh: true,
credentials: map[string]any{
"access_token": "provider-access",
"refresh_token": "provider-refresh",
},
}
result, err := NewOAuthRefreshAPI(repo, cache).RefreshIfNeeded(context.Background(), account, executor, time.Hour)
require.Error(t, err)
require.Nil(t, result)
var containmentErr *providerCycleContainmentRefreshError
require.ErrorAs(t, err, &containmentErr)
require.Equal(t, 2, repo.getByIDCalls)
require.Equal(t, 1, repo.successCASCalls)
require.Equal(t, 1, cache.deleteCalls, "a committed credential rotation must invalidate the pre-rotation access-token cache")
require.NoError(t, cache.deleteCtxErr)
}
func TestRefreshIfNeeded_DBRereadFails(t *testing.T) {
account := &Account{ID: 8, Platform: PlatformAnthropic, Type: AccountTypeOAuth, Status: StatusActive}
repo := &refreshAPIAccountRepo{
@@ -278,21 +476,23 @@ func TestRefreshIfNeeded_DBRereadFails(t *testing.T) {
api := NewOAuthRefreshAPI(repo, cache)
result, err := api.RefreshIfNeeded(context.Background(), account, executor, 3*time.Minute)
require.ErrorContains(t, err, "oauth refresh account reread")
require.Error(t, err)
var stateUnavailable *oauthRefreshStateUnavailableError
require.ErrorAs(t, err, &stateUnavailable)
require.Nil(t, result)
require.Zero(t, executor.refreshCalls, "must not refresh with the stale caller snapshot")
require.Zero(t, executor.refreshCalls, "a failed DB reread must not refresh stale credentials")
require.Zero(t, repo.updateCalls)
require.Equal(t, 1, cache.releaseCalls)
}
func TestRefreshIfNeeded_DBRereadNilFailsClosed(t *testing.T) {
account := &Account{ID: 81, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive}
func TestRefreshIfNeeded_RequestPathDBRereadNilFailsClosed(t *testing.T) {
account := &Account{ID: 81, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true}
repo := &refreshAPIAccountRepo{}
cache := &refreshAPICacheStub{lockResult: true}
executor := &refreshAPIExecutorStub{needsRefresh: true}
api := NewOAuthRefreshAPI(repo, cache)
result, err := api.RefreshIfNeeded(context.Background(), account, executor, 3*time.Minute)
result, err := api.RefreshIfNeeded(withOAuthRefreshRequestPath(context.Background()), account, executor, 3*time.Minute)
require.ErrorIs(t, err, errOAuthRefreshAccountStateChanged)
require.Nil(t, result)
@@ -301,14 +501,14 @@ func TestRefreshIfNeeded_DBRereadNilFailsClosed(t *testing.T) {
require.Equal(t, 1, cache.releaseCalls)
}
func TestRefreshIfNeeded_DBRereadInactiveFailsClosed(t *testing.T) {
account := &Account{ID: 82, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive}
func TestRefreshIfNeeded_RequestPathDBRereadInactiveFailsClosed(t *testing.T) {
account := &Account{ID: 82, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true}
freshAccount := &Account{ID: account.ID, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusDisabled}
repo := &refreshAPIAccountRepo{account: freshAccount}
executor := &refreshAPIExecutorStub{needsRefresh: true}
api := NewOAuthRefreshAPI(repo, nil)
result, err := api.RefreshIfNeeded(context.Background(), account, executor, 3*time.Minute)
result, err := api.RefreshIfNeeded(withOAuthRefreshRequestPath(context.Background()), account, executor, 3*time.Minute)
require.ErrorContains(t, err, "account is not active")
require.Nil(t, result)
@@ -316,7 +516,7 @@ func TestRefreshIfNeeded_DBRereadInactiveFailsClosed(t *testing.T) {
require.Zero(t, repo.updateCalls)
}
func TestRefreshIfNeeded_DBRereadRevalidatesExecutorContract(t *testing.T) {
func TestRefreshIfNeeded_RequestPathDBRereadRevalidatesExecutorContract(t *testing.T) {
tests := []struct {
name string
freshPlatform string
@@ -328,43 +528,130 @@ func TestRefreshIfNeeded_DBRereadRevalidatesExecutorContract(t *testing.T) {
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
account := &Account{ID: 83, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive}
freshAccount := &Account{ID: account.ID, Platform: tt.freshPlatform, Type: tt.freshType, Status: StatusActive}
account := &Account{ID: 83, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true}
freshAccount := &Account{ID: account.ID, Platform: tt.freshPlatform, Type: tt.freshType, Status: StatusActive, Schedulable: true}
repo := &refreshAPIAccountRepo{account: freshAccount}
executor := NewGrokTokenRefresher(nil)
api := NewOAuthRefreshAPI(repo, nil)
result, err := api.RefreshIfNeeded(context.Background(), account, executor, 3*time.Minute)
result, err := api.RefreshIfNeeded(withOAuthRefreshRequestPath(context.Background()), account, executor, 3*time.Minute)
require.ErrorContains(t, err, "no longer refreshable")
require.ErrorIs(t, err, errOAuthRefreshAccountStateChanged)
require.Nil(t, result)
require.Zero(t, repo.updateCalls)
})
}
}
func TestRefreshIfNeeded_DBRereadMissingGrokRefreshCredentialReturnsPermanentSignal(t *testing.T) {
func TestRefreshIfNeeded_LocalLockWaitHonorsContext(t *testing.T) {
account := &Account{ID: 80, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive}
repo := &refreshAPIAccountRepo{account: account}
executor := &refreshAPIExecutorStub{needsRefresh: true}
api := NewOAuthRefreshAPI(repo, nil)
lock := api.getLocalLock(executor.CacheKey(account))
require.NoError(t, lock.Lock(context.Background()))
defer lock.Unlock()
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
defer cancel()
result, err := api.RefreshIfNeeded(ctx, account, executor, time.Hour)
require.ErrorIs(t, err, context.DeadlineExceeded)
require.Nil(t, result)
require.Zero(t, executor.refreshCalls)
}
func TestRefreshIfNeeded_ReleasesDistributedLockAfterParentCancellation(t *testing.T) {
account := &Account{ID: 81, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive}
repo := &refreshAPIAccountRepo{account: account}
cache := &refreshAPICacheStub{lockResult: true}
ctx, cancel := context.WithCancel(context.Background())
executor := &refreshAPIExecutorStub{
needsRefresh: true,
err: errors.New("temporary provider error"),
onRefresh: cancel,
}
api := NewOAuthRefreshAPI(repo, cache)
_, err := api.RefreshIfNeeded(ctx, account, executor, time.Hour)
require.Error(t, err)
require.Equal(t, 1, cache.releaseCalls)
require.NoError(t, cache.releaseCtxErr, "lock cleanup must not reuse the canceled attempt context")
}
func TestRefreshIfNeeded_RevalidatesFreshAccountBeforeRefresh(t *testing.T) {
selected := &Account{ID: 82, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive}
tests := []struct {
name string
fresh *Account
}{
{name: "converted to API key", fresh: &Account{ID: 82, Platform: PlatformGrok, Type: AccountTypeAPIKey, Status: StatusActive}},
{name: "disabled", fresh: &Account{ID: 82, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusDisabled}},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
repo := &refreshAPIAccountRepo{account: tt.fresh}
executor := &refreshAPIExecutorStub{
needsRefresh: true,
canRefresh: func(account *Account) bool {
return account.Platform == PlatformGrok && account.Type == AccountTypeOAuth
},
}
api := NewOAuthRefreshAPI(repo, nil)
result, err := api.RefreshIfNeeded(context.Background(), selected, executor, time.Hour)
require.NoError(t, err)
require.False(t, result.Refreshed)
require.Zero(t, executor.refreshCalls)
require.Zero(t, repo.updateCalls)
})
}
}
func TestRefreshIfNeeded_RequestPathDBRereadMissingGrokRefreshCredentialReturnsPermanentSignal(t *testing.T) {
account := &Account{
ID: 84,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
ID: 84,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Credentials: map[string]any{
"refresh_token": "caller-snapshot-refresh-token",
},
}
freshAccount := &Account{ID: account.ID, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive}
freshAccount := &Account{ID: account.ID, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true}
repo := &refreshAPIAccountRepo{account: freshAccount}
executor := NewGrokTokenRefresher(nil)
api := NewOAuthRefreshAPI(repo, nil)
result, err := api.RefreshIfNeeded(context.Background(), account, executor, 3*time.Minute)
result, err := api.RefreshIfNeeded(withOAuthRefreshRequestPath(context.Background()), account, executor, 3*time.Minute)
require.ErrorIs(t, err, errGrokOAuthRefreshTokenMissing)
require.Nil(t, result)
require.Zero(t, repo.updateCalls)
}
func TestRefreshIfNeeded_LateSuccessAfterDeadlineDoesNotPersist(t *testing.T) {
account := &Account{ID: 85, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive}
repo := &refreshAPIAccountRepo{account: account}
executor := &refreshAPIExecutorStub{
needsRefresh: true,
credentials: map[string]any{"access_token": "late-token"},
delay: 30 * time.Millisecond,
}
api := NewOAuthRefreshAPI(repo, nil)
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
defer cancel()
result, err := api.RefreshIfNeeded(ctx, account, executor, time.Hour)
require.ErrorIs(t, err, context.DeadlineExceeded)
require.Nil(t, result)
require.Zero(t, repo.updateCredentialsCalls, "late credentials must not cross the unified API persistence boundary")
}
func TestRefreshIfNeeded_NilCredentials(t *testing.T) {
account := &Account{ID: 9, Platform: PlatformGemini, Type: AccountTypeOAuth, Status: StatusActive}
repo := &refreshAPIAccountRepo{account: account}
@@ -483,12 +770,12 @@ type refreshAPIAccountRepoWithRace struct {
func (r *refreshAPIAccountRepoWithRace) GetByID(_ context.Context, _ int64) (*Account, error) {
r.getByIDCalls++
if r.getByIDCalls > 1 && r.raceAccount != nil {
return r.raceAccount, nil
return activeRefreshAPITestAccount(r.raceAccount), nil
}
if r.getByIDErr != nil {
return nil, r.getByIDErr
}
return r.account, nil
return activeRefreshAPITestAccount(r.account), nil
}
// ========== Race recovery tests ==========
@@ -554,7 +841,9 @@ func TestRefreshIfNeeded_InvalidGrantGenuine(t *testing.T) {
result, err := api.RefreshIfNeeded(context.Background(), account, executor, 3*time.Minute)
require.Error(t, err, "genuine invalid_grant should propagate error")
require.Nil(t, result)
require.NotNil(t, result)
require.NotNil(t, result.Account)
require.Equal(t, "revoked-rt", result.Account.GetCredential("refresh_token"))
require.Contains(t, err.Error(), "invalid_grant")
}
@@ -580,7 +869,9 @@ func TestRefreshIfNeeded_InvalidGrantDBRereadFailsOnRecovery(t *testing.T) {
result, err := api.RefreshIfNeeded(context.Background(), account, executor, 3*time.Minute)
require.Error(t, err, "should propagate error when recovery DB re-read fails")
require.Nil(t, result)
require.NotNil(t, result)
require.NotNil(t, result.Account)
require.Equal(t, "old-rt", result.Account.GetCredential("refresh_token"))
}
func TestRefreshIfNeeded_LocalMutexSerializesConcurrent(t *testing.T) {
@@ -0,0 +1,949 @@
package service
import (
"context"
"errors"
"fmt"
"reflect"
"sort"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/stretchr/testify/require"
)
type poolHealthAccountRepo struct {
AccountRepository
mu sync.Mutex
pages map[int64][]Account
requests []OAuthRefreshPageOptions
updatedCredentialIDs []int64
setErrorCalls int
setTempUnschedCalls int
getByIDErr error
}
func (r *poolHealthAccountRepo) GetByID(_ context.Context, _ int64) (*Account, error) {
if r.getByIDErr != nil {
return nil, r.getByIDErr
}
return nil, ErrAccountNotFound
}
func (r *poolHealthAccountRepo) ListOAuthRefreshCandidatePage(_ context.Context, options OAuthRefreshPageOptions) (*OAuthRefreshCandidatePage, error) {
r.mu.Lock()
defer r.mu.Unlock()
r.requests = append(r.requests, options)
accounts := append([]Account(nil), r.pages[options.AfterID]...)
page := &OAuthRefreshCandidatePage{Accounts: accounts, HasMore: len(accounts) == options.Limit}
if len(accounts) > 0 {
page.NextAfterID = accounts[len(accounts)-1].ID
}
return page, nil
}
func (r *poolHealthAccountRepo) UpdateCredentials(_ context.Context, id int64, _ map[string]any) error {
r.mu.Lock()
defer r.mu.Unlock()
r.updatedCredentialIDs = append(r.updatedCredentialIDs, id)
return nil
}
func (r *poolHealthAccountRepo) UpdateGrokOAuthCredentialsIfUnchanged(
_ context.Context,
id int64,
_ map[string]any,
_ *int64,
_ map[string]any,
) (bool, error) {
r.mu.Lock()
defer r.mu.Unlock()
r.updatedCredentialIDs = append(r.updatedCredentialIDs, id)
return true, nil
}
func (r *poolHealthAccountRepo) SetError(context.Context, int64, string) error {
r.mu.Lock()
defer r.mu.Unlock()
r.setErrorCalls++
return nil
}
func (r *poolHealthAccountRepo) SetGrokOAuthErrorIfCredentialsUnchanged(context.Context, int64, map[string]any, string) (bool, error) {
return false, nil
}
func (r *poolHealthAccountRepo) SetGrokOAuthRefreshErrorIfCredentialsUnchanged(context.Context, int64, map[string]any, *int64, string) (bool, error) {
r.mu.Lock()
defer r.mu.Unlock()
r.setErrorCalls++
return true, nil
}
func (r *poolHealthAccountRepo) SetGrokOAuthRefreshTempUnschedulableIfCredentialsUnchanged(context.Context, int64, map[string]any, *int64, time.Time, string) (bool, error) {
r.mu.Lock()
defer r.mu.Unlock()
r.setTempUnschedCalls++
return true, nil
}
func (r *poolHealthAccountRepo) SetTempUnschedulable(context.Context, int64, time.Time, string) error {
r.mu.Lock()
defer r.mu.Unlock()
r.setTempUnschedCalls++
return nil
}
func (r *poolHealthAccountRepo) snapshot() ([]OAuthRefreshPageOptions, []int64, int, int) {
r.mu.Lock()
defer r.mu.Unlock()
return append([]OAuthRefreshPageOptions(nil), r.requests...), append([]int64(nil), r.updatedCredentialIDs...), r.setErrorCalls, r.setTempUnschedCalls
}
type poolHealthRefresher struct {
err error
delay time.Duration
startDelays []time.Duration
ignoreContext bool
cancel context.CancelFunc
newCredentials map[string]any
calls atomic.Int64
active atomic.Int64
maxActive atomic.Int64
startMu sync.Mutex
startTimes []time.Time
}
type countingRefreshAttemptGate struct {
calls atomic.Int64
}
type rejectedRefreshAttemptGate struct {
err error
}
type poolHealthTokenCacheStub struct {
GeminiTokenCache
}
type tripBeforeRateAdmissionGate struct {
state *tokenRefreshProviderState
}
func (g *tripBeforeRateAdmissionGate) acquire(ctx context.Context) (func(), error) {
release, err := g.state.acquire(ctx)
if err != nil {
return nil, err
}
g.state.mu.Lock()
g.state.tripped = true
g.state.mu.Unlock()
return release, nil
}
func (g *tripBeforeRateAdmissionGate) acquireRate(ctx context.Context) (func(), error) {
return g.state.acquireRate(ctx)
}
type breakerTripAccountRepo struct {
*productionPathRateRepo
setErrorCalls atomic.Int64
setTempCalls atomic.Int64
}
func (r *breakerTripAccountRepo) SetGrokOAuthRefreshErrorIfCredentialsUnchanged(context.Context, int64, map[string]any, *int64, string) (bool, error) {
r.setErrorCalls.Add(1)
return true, nil
}
func (r *breakerTripAccountRepo) SetGrokOAuthRefreshTempUnschedulableIfCredentialsUnchanged(context.Context, int64, map[string]any, *int64, time.Time, string) (bool, error) {
r.setTempCalls.Add(1)
return true, nil
}
func (g *rejectedRefreshAttemptGate) acquire(context.Context) (func(), error) {
return nil, g.err
}
type productionPathRateRepo struct {
AccountRepository
mu sync.Mutex
accounts map[int64]*Account
}
func (r *productionPathRateRepo) GetByID(_ context.Context, id int64) (*Account, error) {
r.mu.Lock()
defer r.mu.Unlock()
account := r.accounts[id]
if account == nil {
return nil, ErrAccountNotFound
}
return snapshotOAuthRefreshAccount(account), nil
}
func (r *productionPathRateRepo) UpdateCredentials(_ context.Context, id int64, credentials map[string]any) error {
r.mu.Lock()
defer r.mu.Unlock()
account := r.accounts[id]
if account == nil {
return ErrAccountNotFound
}
account.Credentials = shallowCopyMap(credentials)
return nil
}
func (r *productionPathRateRepo) UpdateGrokOAuthCredentialsIfUnchanged(
_ context.Context,
id int64,
expectedCredentials map[string]any,
expectedProxyID *int64,
credentials map[string]any,
) (bool, error) {
r.mu.Lock()
defer r.mu.Unlock()
account := r.accounts[id]
if account == nil || !reflect.DeepEqual(account.Credentials, expectedCredentials) ||
!reflect.DeepEqual(account.ProxyID, expectedProxyID) {
return false, nil
}
account.Credentials = shallowCopyMap(credentials)
return true, nil
}
type productionPathRefreshStart struct {
accountID int64
at time.Time
}
type productionPathRateExecutor struct {
firstStarted chan struct{}
releaseFirst chan struct{}
calls atomic.Int64
startMu sync.Mutex
starts []productionPathRefreshStart
}
func (e *productionPathRateExecutor) CacheKey(account *Account) string {
return fmt.Sprintf("production-path-rate:%d", account.ID)
}
func (e *productionPathRateExecutor) CanRefresh(account *Account) bool {
return account != nil && account.IsGrokOAuth()
}
func (e *productionPathRateExecutor) NeedsRefresh(account *Account, _ time.Duration) bool {
needsRefresh, _ := account.Credentials["needs_refresh"].(bool)
return needsRefresh
}
func (e *productionPathRateExecutor) Refresh(ctx context.Context, account *Account) (map[string]any, error) {
call := e.calls.Add(1)
e.startMu.Lock()
e.starts = append(e.starts, productionPathRefreshStart{accountID: account.ID, at: time.Now()})
e.startMu.Unlock()
if call == 1 {
close(e.firstStarted)
select {
case <-ctx.Done():
return nil, ctx.Err()
case <-e.releaseFirst:
}
}
return map[string]any{
"access_token": fmt.Sprintf("fresh-access-%d", account.ID),
"refresh_token": fmt.Sprintf("fresh-refresh-%d", account.ID),
"needs_refresh": false,
}, nil
}
func (e *productionPathRateExecutor) startsSnapshot() []productionPathRefreshStart {
e.startMu.Lock()
defer e.startMu.Unlock()
return append([]productionPathRefreshStart(nil), e.starts...)
}
func (g *countingRefreshAttemptGate) acquire(ctx context.Context) (func(), error) {
if err := ctx.Err(); err != nil {
return nil, err
}
g.calls.Add(1)
return func() {}, nil
}
func (r *poolHealthRefresher) CacheKey(account *Account) string {
return fmt.Sprintf("pool-health:%d", account.ID)
}
func (r *poolHealthRefresher) CanRefresh(account *Account) bool {
return account != nil && account.Platform == PlatformGrok && account.Type == AccountTypeOAuth
}
func (r *poolHealthRefresher) NeedsRefresh(*Account, time.Duration) bool { return true }
func (r *poolHealthRefresher) Refresh(ctx context.Context, _ *Account) (map[string]any, error) {
r.calls.Add(1)
active := r.active.Add(1)
defer r.active.Add(-1)
r.startMu.Lock()
startIndex := len(r.startTimes)
r.startTimes = append(r.startTimes, time.Now())
delay := r.delay
if startIndex < len(r.startDelays) {
delay = r.startDelays[startIndex]
}
r.startMu.Unlock()
for {
maxActive := r.maxActive.Load()
if active <= maxActive || r.maxActive.CompareAndSwap(maxActive, active) {
break
}
}
if r.cancel != nil {
r.cancel()
}
if delay > 0 {
if r.ignoreContext {
time.Sleep(delay)
} else {
timer := time.NewTimer(delay)
defer timer.Stop()
select {
case <-ctx.Done():
return nil, ctx.Err()
case <-timer.C:
}
}
}
if r.err != nil {
return nil, r.err
}
if r.newCredentials != nil {
credentials := make(map[string]any, len(r.newCredentials))
for key, value := range r.newCredentials {
credentials[key] = value
}
return credentials, nil
}
return map[string]any{"access_token": "new-token", "refresh_token": "new-refresh-token"}, nil
}
func (r *poolHealthRefresher) startsSnapshot() []time.Time {
r.startMu.Lock()
defer r.startMu.Unlock()
return append([]time.Time(nil), r.startTimes...)
}
func grokPoolAccount(id int64) Account {
return Account{
ID: id,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Credentials: map[string]any{
"access_token": "old-token",
"refresh_token": "refresh-token",
},
}
}
func newPoolHealthService(repo *poolHealthAccountRepo, refresher *poolHealthRefresher, cfg config.TokenRefreshConfig) *TokenRefreshService {
return &TokenRefreshService{
accountRepo: repo,
candidatePager: repo,
registrations: []tokenRefreshRegistration{{
platform: PlatformGrok,
refresher: refresher,
executor: refresher,
}},
refreshPolicy: DefaultBackgroundRefreshPolicy(),
cfg: &cfg,
}
}
func TestTokenRefreshService_RegistrationsAreCandidateEligibilitySource(t *testing.T) {
cfg := &config.Config{}
svc := NewTokenRefreshService(nil, nil, nil, nil, nil, nil, nil, cfg, nil)
require.Equal(t, []string{
PlatformAnthropic,
PlatformOpenAI,
PlatformGemini,
PlatformAntigravity,
PlatformGrok,
}, svc.eligiblePlatforms())
require.Len(t, svc.registrations, 5)
for _, registration := range svc.registrations {
require.NotNil(t, registration.refresher)
require.NotNil(t, registration.executor)
}
}
func TestTokenRefreshService_ProcessRefreshPagesByStableCursor(t *testing.T) {
repo := &poolHealthAccountRepo{pages: map[int64][]Account{
0: {grokPoolAccount(1), grokPoolAccount(2)},
2: {grokPoolAccount(3)},
}}
refresher := &poolHealthRefresher{}
svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{
RefreshBeforeExpiryHours: 1,
MaxRetries: 1,
CandidatePageSize: 2,
ProviderConcurrency: 4,
ProviderQPS: 10000,
AttemptTimeoutSeconds: 1,
CycleTimeoutSeconds: 2,
})
svc.processRefreshContext(context.Background())
requests, updatedIDs, _, _ := repo.snapshot()
require.Len(t, requests, 2)
require.Equal(t, int64(0), requests[0].AfterID)
require.Equal(t, int64(2), requests[1].AfterID)
require.Equal(t, []string{PlatformGrok}, requests[0].Platforms)
require.True(t, requests[0].ActiveOnly)
require.True(t, requests[0].RequireRefreshToken)
require.True(t, requests[0].ExcludeRetryCooldown)
sort.Slice(updatedIDs, func(i, j int) bool { return updatedIDs[i] < updatedIDs[j] })
require.Equal(t, []int64{1, 2, 3}, updatedIDs)
require.Zero(t, svc.candidateAfterID(), "a short final page must wrap the next cycle to the beginning")
}
func TestTokenRefreshService_BoundsPerProviderConcurrency(t *testing.T) {
accounts := make([]Account, 0, 8)
for id := int64(1); id <= 8; id++ {
accounts = append(accounts, grokPoolAccount(id))
}
repo := &poolHealthAccountRepo{pages: map[int64][]Account{0: accounts}}
refresher := &poolHealthRefresher{delay: 20 * time.Millisecond}
svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{
MaxRetries: 1,
CandidatePageSize: 20,
ProviderConcurrency: 2,
ProviderQPS: 10000,
AttemptTimeoutSeconds: 1,
CycleTimeoutSeconds: 2,
})
svc.processRefreshContext(context.Background())
require.Equal(t, int64(8), refresher.calls.Load())
require.Equal(t, int64(2), refresher.maxActive.Load())
}
func TestTokenRefreshRateGate_ReservesSpacedSlotsAndHonorsCancellation(t *testing.T) {
const interval = 25 * time.Millisecond
gate := newTokenRefreshRateGateWithInterval(interval)
base := time.Unix(1_700_000_000, 0)
require.Equal(t, base, gate.reserveSlot(base))
require.Equal(t, base.Add(interval), gate.reserveSlot(base))
require.Equal(t, base.Add(2*interval), gate.reserveSlot(base))
jumped := base.Add(time.Second)
require.Equal(t, jumped, gate.reserveSlot(jumped), "an idle gate should not retain stale delay")
cancelGate := newTokenRefreshRateGateWithInterval(time.Hour)
require.NoError(t, cancelGate.wait(context.Background()), "the first slot is immediately available")
ctx, cancel := context.WithCancel(context.Background())
cancel()
started := time.Now()
require.ErrorIs(t, cancelGate.wait(ctx), context.Canceled)
require.Less(t, time.Since(started), 100*time.Millisecond, "cancellation must not wait for the reserved slot")
}
func TestTokenRefreshService_RetriesAcquireRateSlotPerAttempt(t *testing.T) {
repo := &poolHealthAccountRepo{}
refresher := &poolHealthRefresher{err: errors.New("temporary provider failure")}
svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{MaxRetries: 3})
gate := &countingRefreshAttemptGate{}
account := grokPoolAccount(44)
err := svc.refreshWithRetryWithRateGate(context.Background(), &account, refresher, nil, time.Hour, gate)
require.Error(t, err)
require.Equal(t, int64(3), refresher.calls.Load())
require.Equal(t, int64(3), gate.calls.Load(), "every upstream retry must consume a provider rate slot")
}
func TestTokenRefreshService_ProcessProviderAccountsLegacyNilReleaseGateIsSafe(t *testing.T) {
repo := &poolHealthAccountRepo{}
refresher := &poolHealthRefresher{}
svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{
MaxRetries: 1,
ProviderConcurrency: 1,
})
state := &tokenRefreshProviderState{
service: svc,
registration: tokenRefreshRegistration{
platform: PlatformGrok,
refresher: refresher,
// nil executor deliberately exercises the legacy/direct fallback.
executor: nil,
},
// Admission rejection validly returns no release callback. The direct
// fallback must propagate the skip without dereferencing that nil handle.
rateGate: &rejectedRefreshAttemptGate{err: errRefreshSkipped},
poolGate: nil,
}
account := grokPoolAccount(45)
refreshed, skipped, failed := svc.processProviderAccounts(
context.Background(),
state,
[]*Account{&account},
time.Hour,
)
require.Zero(t, refreshed)
require.Equal(t, 1, skipped)
require.Zero(t, failed)
require.Zero(t, refresher.calls.Load(), "rejected rate admission must not reach the legacy upstream refresher")
}
func TestTokenRefreshService_ProviderRateGateIsSharedAcrossRuns(t *testing.T) {
svc := &TokenRefreshService{cfg: &config.TokenRefreshConfig{ProviderQPS: 40}}
first := svc.providerRateGate(PlatformGrok)
second := svc.providerRateGate(PlatformGrok)
require.Same(t, first, second, "background cycles and reconciliation must share the process-local provider limiter")
base := time.Unix(1_700_000_000, 0)
require.Equal(t, base, first.reserveSlot(base))
require.Equal(t, base.Add(25*time.Millisecond), second.reserveSlot(base))
}
func TestTokenRefreshService_ProviderConcurrencyGateIsSharedAcrossBackgroundAndConcurrentAdminReconciliation(t *testing.T) {
accounts := []Account{
grokPoolAccount(1),
grokPoolAccount(2),
grokPoolAccount(3),
grokPoolAccount(4),
}
repo := &poolHealthAccountRepo{pages: map[int64][]Account{0: accounts}}
refresher := &poolHealthRefresher{delay: 80 * time.Millisecond}
svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{
RefreshBeforeExpiryHours: 1,
MaxRetries: 1,
CandidatePageSize: 20,
ProviderConcurrency: 2,
ProviderQPS: 100,
ProviderFailureThreshold: 20,
AttemptTimeoutSeconds: 1,
CycleTimeoutSeconds: 3,
})
firstGate := svc.providerConcurrencyGate(PlatformGrok)
require.Same(t, firstGate, svc.providerConcurrencyGate(PlatformGrok))
start := make(chan struct{})
adminErrors := make(chan error, 2)
var wg sync.WaitGroup
wg.Add(3)
go func() {
defer wg.Done()
<-start
svc.processRefreshContext(context.Background())
}()
for i := 0; i < 2; i++ {
go func() {
defer wg.Done()
<-start
_, err := svc.ReconcileGrokOAuth(context.Background(), GrokOAuthReconcileInput{Apply: true, Limit: 20})
adminErrors <- err
}()
}
close(start)
wg.Wait()
close(adminErrors)
for err := range adminErrors {
require.NoError(t, err)
}
require.Equal(t, int64(12), refresher.calls.Load(), "background and both admin calls must all execute")
require.Equal(t, int64(2), refresher.maxActive.Load(),
"all entry points must share the configured per-provider upstream concurrency cap")
}
func TestTokenRefreshService_SaturatedProviderPreservesConcurrencyAndActualQPSStartSpacing(t *testing.T) {
const (
providerConcurrency = 2
providerQPS = 20
attemptCount = 8
)
repo := &poolHealthAccountRepo{}
refresher := &poolHealthRefresher{
// The first two QPS-spaced attempts finish together. If queued callers
// reserve QPS slots before acquiring provider capacity, two expired
// reservations can then burst upstream at the same time.
startDelays: []time.Duration{
220 * time.Millisecond,
170 * time.Millisecond,
20 * time.Millisecond,
20 * time.Millisecond,
20 * time.Millisecond,
20 * time.Millisecond,
20 * time.Millisecond,
20 * time.Millisecond,
},
}
svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{
MaxRetries: 1,
ProviderConcurrency: providerConcurrency,
ProviderQPS: providerQPS,
AttemptTimeoutSeconds: 1,
})
registration := svc.registrations[0]
sharedRateGate := svc.providerRateGate(PlatformGrok)
sharedPoolGate := svc.providerConcurrencyGate(PlatformGrok)
start := make(chan struct{})
errorsCh := make(chan error, attemptCount)
var wg sync.WaitGroup
for i := 0; i < attemptCount; i++ {
account := grokPoolAccount(int64(i + 1))
state := &tokenRefreshProviderState{
service: svc,
registration: registration,
rateGate: sharedRateGate,
poolGate: sharedPoolGate,
}
wg.Add(1)
go func() {
defer wg.Done()
<-start
errorsCh <- svc.refreshWithRetryWithRateGate(context.Background(), &account, refresher, nil, time.Hour, state)
}()
}
close(start)
wg.Wait()
close(errorsCh)
for err := range errorsCh {
require.NoError(t, err)
}
require.Equal(t, int64(providerConcurrency), refresher.maxActive.Load(),
"the scripted attempts must actually saturate the provider semaphore")
starts := refresher.startsSnapshot()
require.Len(t, starts, attemptCount)
configuredSpacing := time.Second / time.Duration(providerQPS)
minimumObservedSpacing := configuredSpacing - 10*time.Millisecond
actualMinimumSpacing := starts[1].Sub(starts[0])
for i := 1; i < len(starts); i++ {
spacing := starts[i].Sub(starts[i-1])
if spacing < actualMinimumSpacing {
actualMinimumSpacing = spacing
}
require.GreaterOrEqualf(t, spacing, minimumObservedSpacing,
"upstream starts %d and %d violated configured QPS spacing", i-1, i)
}
t.Logf("max_active=%d configured_concurrency=%d minimum_start_spacing=%s configured_spacing=%s",
refresher.maxActive.Load(), providerConcurrency, actualMinimumSpacing, configuredSpacing)
}
func TestTokenRefreshService_ProductionPathRatesOnlyActualRefreshAfterSameAccountContention(t *testing.T) {
const interval = 200 * time.Millisecond
accountOne := grokPoolAccount(71)
accountOne.Credentials["needs_refresh"] = true
accountTwo := grokPoolAccount(72)
accountTwo.Credentials["needs_refresh"] = true
firstSelection := snapshotOAuthRefreshAccount(&accountOne)
contendingSelection := snapshotOAuthRefreshAccount(&accountOne)
differentSelection := snapshotOAuthRefreshAccount(&accountTwo)
repo := &productionPathRateRepo{accounts: map[int64]*Account{
accountOne.ID: snapshotOAuthRefreshAccount(&accountOne),
accountTwo.ID: snapshotOAuthRefreshAccount(&accountTwo),
}}
executor := &productionPathRateExecutor{
firstStarted: make(chan struct{}),
releaseFirst: make(chan struct{}),
}
svc := &TokenRefreshService{
accountRepo: repo,
refreshAPI: NewOAuthRefreshAPI(repo, nil),
refreshPolicy: DefaultBackgroundRefreshPolicy(),
cfg: &config.TokenRefreshConfig{MaxRetries: 1},
attemptTimeoutOverride: 2 * time.Second,
}
state := &tokenRefreshProviderState{
service: svc,
rateGate: newTokenRefreshRateGateWithInterval(interval),
poolGate: newTokenRefreshConcurrencyGate(2),
}
errorsCh := make(chan error, 3)
go func() {
errorsCh <- svc.refreshWithRetryWithRateGate(context.Background(), firstSelection, executor, executor, time.Hour, state)
}()
select {
case <-executor.firstStarted:
case <-time.After(time.Second):
require.FailNow(t, "first production-path refresh did not reach the upstream executor")
}
go func() {
errorsCh <- svc.refreshWithRetryWithRateGate(context.Background(), contendingSelection, executor, executor, time.Hour, state)
}()
require.Eventually(t, func() bool {
return len(state.poolGate.slots) == 2
}, time.Second, time.Millisecond, "same-account contender must hold the second provider slot while waiting on the local refresh lock")
go func() {
errorsCh <- svc.refreshWithRetryWithRateGate(context.Background(), differentSelection, executor, executor, time.Hour, state)
}()
close(executor.releaseFirst)
skipped := 0
for i := 0; i < 3; i++ {
err := <-errorsCh
if errors.Is(err, errRefreshSkipped) {
skipped++
continue
}
require.NoError(t, err)
}
require.Equal(t, 1, skipped, "the same-account contender must reread the refreshed row and skip without upstream admission")
starts := executor.startsSnapshot()
require.Len(t, starts, 2, "only the two accounts that actually refresh may consume QPS admission")
require.Equal(t, int64(71), starts[0].accountID)
require.Equal(t, int64(72), starts[1].accountID)
spacing := starts[1].at.Sub(starts[0].at)
require.GreaterOrEqual(t, spacing, interval-30*time.Millisecond)
require.Less(t, spacing, 350*time.Millisecond,
"a same-account lock waiter must not consume a rate slot and push the different-account refresh to the second interval")
t.Logf("actual_refresh_calls=%d actual_start_spacing=%s configured_spacing=%s", executor.calls.Load(), spacing, interval)
}
func TestTokenRefreshService_ProviderTripBeforeRateAdmissionSkipsWithoutAccountMutation(t *testing.T) {
account := grokPoolAccount(73)
stored := snapshotOAuthRefreshAccount(&account)
repo := &breakerTripAccountRepo{productionPathRateRepo: &productionPathRateRepo{
accounts: map[int64]*Account{account.ID: stored},
}}
refresher := &poolHealthRefresher{}
svc := &TokenRefreshService{
accountRepo: repo,
refreshAPI: NewOAuthRefreshAPI(repo, nil),
refreshPolicy: DefaultBackgroundRefreshPolicy(),
cfg: &config.TokenRefreshConfig{MaxRetries: 1},
}
state := &tokenRefreshProviderState{
service: svc,
rateGate: newTokenRefreshRateGate(1),
poolGate: newTokenRefreshConcurrencyGate(1),
}
gate := &tripBeforeRateAdmissionGate{state: state}
err := svc.refreshWithRetryWithRateGate(context.Background(), &account, refresher, refresher, time.Hour, gate)
require.ErrorIs(t, err, errRefreshSkipped)
require.Zero(t, refresher.calls.Load(), "a tripped provider must not reach upstream rate admission")
require.Zero(t, repo.setErrorCalls.Load())
require.Zero(t, repo.setTempCalls.Load(), "provider skip must never fall through to per-account cooldown")
}
func TestTokenRefreshService_ConfigBounds(t *testing.T) {
maxInt := int(^uint(0) >> 1)
svc := &TokenRefreshService{cfg: &config.TokenRefreshConfig{
MaxRetries: maxInt,
RetryBackoffSeconds: maxInt,
ProviderFailureThreshold: maxInt,
AttemptTimeoutSeconds: maxInt,
CycleTimeoutSeconds: maxInt,
}}
require.Equal(t, maxTokenRefreshMaxRetries, svc.maxRetries())
require.Equal(t, maxTokenRefreshProviderFailureThreshold, svc.providerFailureThreshold())
require.Equal(t, maxTokenRefreshAttemptTimeout, svc.attemptTimeout())
require.Equal(t, maxTokenRefreshCycleTimeout, svc.cycleTimeout())
require.LessOrEqual(t, svc.retryBackoff(1, maxTokenRefreshMaxRetries), maxTokenRefreshRetryBackoff)
require.Equal(t, maxGrokOAuthReconcilePageSize, svc.grokOAuthReconcileMaxPageSize())
}
func TestTokenRefreshService_AttemptTimeoutStaysInsideDistributedLockLease(t *testing.T) {
cache := &poolHealthTokenCacheStub{}
svc := &TokenRefreshService{
cfg: &config.TokenRefreshConfig{AttemptTimeoutSeconds: int(maxTokenRefreshAttemptTimeout / time.Second)},
refreshAPI: NewOAuthRefreshAPI(&poolHealthAccountRepo{}, cache),
}
require.Equal(t, 55*time.Second, svc.attemptTimeout())
require.Less(t, svc.attemptTimeout(), defaultRefreshLockTTL)
}
func TestTokenRefreshService_SharedProviderFailureContainsCycleWithoutAccountMutation(t *testing.T) {
accounts := make([]Account, 0, 5)
for id := int64(1); id <= 5; id++ {
accounts = append(accounts, grokPoolAccount(id))
}
repo := &poolHealthAccountRepo{pages: map[int64][]Account{0: accounts}}
refresher := &poolHealthRefresher{err: errors.New("invalid_client: provider configuration rejected")}
svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{
MaxRetries: 1,
CandidatePageSize: 10,
ProviderConcurrency: 4,
ProviderQPS: 10000,
ProviderFailureThreshold: 3,
AttemptTimeoutSeconds: 1,
CycleTimeoutSeconds: 2,
})
svc.processRefreshContext(context.Background())
_, _, setErrorCalls, setTempUnschedCalls := repo.snapshot()
require.Equal(t, int64(1), refresher.calls.Load(), "shared provider configuration failures must open the in-cycle breaker immediately")
require.Zero(t, setErrorCalls, "shared provider failures must not mass-disable accounts")
require.Zero(t, setTempUnschedCalls, "shared provider failures must not mutate per-account scheduling state")
}
func TestTokenRefreshService_SharedDBRereadFailureContainsCycleWithoutAccountMutation(t *testing.T) {
accounts := []Account{grokPoolAccount(1), grokPoolAccount(2), grokPoolAccount(3)}
repo := &poolHealthAccountRepo{
pages: map[int64][]Account{0: accounts},
getByIDErr: errors.New("database unavailable"),
}
refresher := &poolHealthRefresher{}
svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{
MaxRetries: 3,
CandidatePageSize: 10,
ProviderConcurrency: 4,
ProviderQPS: 10000,
AttemptTimeoutSeconds: 1,
CycleTimeoutSeconds: 2,
})
svc.refreshAPI = NewOAuthRefreshAPI(repo, nil)
svc.processRefreshContext(context.Background())
_, _, setErrorCalls, setTempUnschedCalls := repo.snapshot()
require.Zero(t, refresher.calls.Load(), "refresh must fail closed before using stale account credentials")
require.Zero(t, setErrorCalls)
require.Zero(t, setTempUnschedCalls, "a shared DB outage must not mutate the selected account")
}
func TestTokenRefreshService_GenericGrokForbiddenContainsCycleWithoutAccountMutation(t *testing.T) {
accounts := []Account{grokPoolAccount(1), grokPoolAccount(2), grokPoolAccount(3)}
repo := &poolHealthAccountRepo{pages: map[int64][]Account{0: accounts}}
refresher := &poolHealthRefresher{err: errors.New(`GROK_OAUTH_ENTITLEMENT_DENIED: token refresh failed: status 403, body: <html>request blocked</html>`)}
svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{
MaxRetries: 1,
CandidatePageSize: 10,
ProviderConcurrency: 4,
ProviderQPS: 10000,
ProviderFailureThreshold: 3,
AttemptTimeoutSeconds: 1,
CycleTimeoutSeconds: 2,
})
svc.processRefreshContext(context.Background())
_, _, setErrorCalls, setTempUnschedCalls := repo.snapshot()
require.Equal(t, int64(1), refresher.calls.Load(), "an ambiguous Grok 403 must contain the provider immediately")
require.Zero(t, setErrorCalls, "a generic 403 is not evidence that an account credential is permanently invalid")
require.Zero(t, setTempUnschedCalls, "provider containment must not mutate account scheduling state")
}
func TestTokenRefreshService_ExplicitGrokEntitlementDenialIsPermanent(t *testing.T) {
repo := &poolHealthAccountRepo{pages: map[int64][]Account{0: {grokPoolAccount(1)}}}
refresher := &poolHealthRefresher{err: errors.New(`GROK_OAUTH_ENTITLEMENT_DENIED: token refresh failed: status 403, body: {"error":"subscription required"}`)}
svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{
MaxRetries: 1,
CandidatePageSize: 10,
ProviderConcurrency: 1,
ProviderQPS: 10000,
AttemptTimeoutSeconds: 1,
CycleTimeoutSeconds: 2,
})
svc.processRefreshContext(context.Background())
_, _, setErrorCalls, setTempUnschedCalls := repo.snapshot()
require.Equal(t, int64(1), refresher.calls.Load())
require.Equal(t, 1, setErrorCalls, "explicit entitlement evidence is an account-permanent failure")
require.Zero(t, setTempUnschedCalls)
}
func TestTokenRefreshService_AttemptTimeoutTripsRetryableProviderThreshold(t *testing.T) {
accounts := []Account{grokPoolAccount(1), grokPoolAccount(2), grokPoolAccount(3)}
repo := &poolHealthAccountRepo{pages: map[int64][]Account{0: accounts}}
refresher := &poolHealthRefresher{delay: time.Second}
svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{
MaxRetries: 1,
CandidatePageSize: 10,
ProviderConcurrency: 1,
ProviderQPS: 10000,
ProviderFailureThreshold: 2,
CycleTimeoutSeconds: 2,
})
svc.attemptTimeoutOverride = 20 * time.Millisecond
svc.processRefreshContext(context.Background())
_, _, setErrorCalls, setTempUnschedCalls := repo.snapshot()
require.Equal(t, int64(2), refresher.calls.Load(), "two attempt timeouts should trip the retryable provider threshold")
require.Zero(t, setErrorCalls)
require.Equal(t, 2, setTempUnschedCalls, "attempt timeouts remain account-transient failures before containment opens")
}
func TestTokenRefreshService_ParentCancellationStopsRetryWithoutAccountMutation(t *testing.T) {
repo := &poolHealthAccountRepo{}
ctx, cancel := context.WithCancel(context.Background())
refresher := &poolHealthRefresher{
err: errors.New("temporary provider failure"),
cancel: cancel,
}
svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{
MaxRetries: 3,
RetryBackoffSeconds: 1,
AttemptTimeoutSeconds: 1,
})
account := grokPoolAccount(42)
err := svc.refreshWithRetry(ctx, &account, refresher, nil, time.Hour)
require.ErrorIs(t, err, context.Canceled)
_, _, setErrorCalls, setTempUnschedCalls := repo.snapshot()
require.Zero(t, setErrorCalls)
require.Zero(t, setTempUnschedCalls)
}
func TestTokenRefreshService_LateSuccessPastAttemptDeadlineIsRejected(t *testing.T) {
repo := &poolHealthAccountRepo{}
refresher := &poolHealthRefresher{
delay: 30 * time.Millisecond,
ignoreContext: true,
}
svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{MaxRetries: 1})
svc.attemptTimeoutOverride = 10 * time.Millisecond
account := grokPoolAccount(43)
err := svc.refreshWithRetry(context.Background(), &account, refresher, nil, time.Hour)
var timeoutErr *refreshAttemptTimeoutError
require.ErrorAs(t, err, &timeoutErr)
_, updatedIDs, setErrorCalls, setTempUnschedCalls := repo.snapshot()
require.Empty(t, updatedIDs, "credentials returned after the deadline must not be persisted")
require.Zero(t, setErrorCalls)
require.Equal(t, 1, setTempUnschedCalls)
}
func TestTokenRefreshService_NonRetryableGrokFailureInvalidatesTokenCache(t *testing.T) {
repo := &poolHealthAccountRepo{}
invalidator := &reconcileInvalidator{}
refresher := &poolHealthRefresher{err: errors.New("invalid_grant: revoked")}
svc := newPoolHealthService(repo, refresher, config.TokenRefreshConfig{MaxRetries: 1})
svc.cacheInvalidator = invalidator
account := grokPoolAccount(77)
err := svc.refreshWithRetry(context.Background(), &account, refresher, nil, time.Hour)
require.Error(t, err)
_, _, setErrorCalls, setTempUnschedCalls := repo.snapshot()
require.Equal(t, 1, setErrorCalls)
require.Zero(t, setTempUnschedCalls)
require.Equal(t, 1, invalidator.count())
}
File diff suppressed because it is too large Load Diff
@@ -4,6 +4,7 @@ import (
"context"
"errors"
"strings"
"sync"
"testing"
"time"
@@ -14,6 +15,7 @@ import (
type tokenRefreshCandidateRepo struct {
AccountRepository
mu sync.Mutex
accounts []Account
updatedCredentialIDs []int64
setErrorCalls int
@@ -24,60 +26,78 @@ type tokenRefreshCandidateRepo struct {
}
func (r *tokenRefreshCandidateRepo) ListActive(context.Context) ([]Account, error) {
r.mu.Lock()
defer r.mu.Unlock()
r.listActiveCalls++
return r.accounts, nil
}
func (r *tokenRefreshCandidateRepo) ListOAuthRefreshCandidates(context.Context) ([]Account, error) {
func (r *tokenRefreshCandidateRepo) ListOAuthRefreshCandidatePage(_ context.Context, options OAuthRefreshPageOptions) (*OAuthRefreshCandidatePage, error) {
candidates := make([]Account, 0, len(r.accounts))
now := time.Now()
for _, account := range r.accounts {
if account.ID <= options.AfterID {
continue
}
refreshToken, _ := account.Credentials["refresh_token"].(string)
inRetryCooldown := account.TempUnschedulableUntil != nil &&
account.TempUnschedulableUntil.After(now) &&
strings.HasPrefix(account.TempUnschedulableReason, "token refresh retry exhausted:")
if account.Status != StatusActive ||
platformAllowed := false
for _, platform := range options.Platforms {
if account.Platform == platform {
platformAllowed = true
break
}
}
if options.ActiveOnly && account.Status != StatusActive ||
account.Type != AccountTypeOAuth ||
!isOAuthRefreshPlatform(account.Platform) ||
strings.TrimSpace(refreshToken) == "" ||
inRetryCooldown {
!platformAllowed ||
options.RequireRefreshToken && strings.TrimSpace(refreshToken) == "" ||
options.ExcludeRetryCooldown && inRetryCooldown {
continue
}
candidates = append(candidates, account)
if len(candidates) == options.Limit {
break
}
}
return candidates, nil
page := &OAuthRefreshCandidatePage{Accounts: candidates, HasMore: len(candidates) == options.Limit}
if len(candidates) > 0 {
page.NextAfterID = candidates[len(candidates)-1].ID
}
return page, nil
}
func (r *tokenRefreshCandidateRepo) UpdateCredentials(_ context.Context, id int64, _ map[string]any) error {
r.mu.Lock()
defer r.mu.Unlock()
r.updatedCredentialIDs = append(r.updatedCredentialIDs, id)
return nil
}
func (r *tokenRefreshCandidateRepo) SetError(context.Context, int64, string) error {
r.mu.Lock()
defer r.mu.Unlock()
r.setErrorCalls++
return nil
}
func (r *tokenRefreshCandidateRepo) SetTempUnschedulable(_ context.Context, _ int64, _ time.Time, reason string) error {
r.mu.Lock()
defer r.mu.Unlock()
r.setTempUnschedCalls++
r.lastTempUnschedReason = reason
return nil
}
func (r *tokenRefreshCandidateRepo) ClearTempUnschedulable(context.Context, int64) error {
r.mu.Lock()
defer r.mu.Unlock()
r.clearTempCalls++
return nil
}
func isOAuthRefreshPlatform(platform string) bool {
switch platform {
case PlatformAnthropic, PlatformOpenAI, PlatformGemini, PlatformAntigravity:
return true
default:
return false
}
}
type tokenRefreshTestRefresher struct {
err error
}
@@ -147,8 +167,13 @@ func TestTokenRefreshService_ProcessRefreshUsesOAuthRefreshCandidates(t *testing
},
}
svc := &TokenRefreshService{
accountRepo: repo,
refreshers: []TokenRefresher{&tokenRefreshTestRefresher{}},
accountRepo: repo,
candidatePager: repo,
registrations: []tokenRefreshRegistration{
{platform: PlatformOpenAI, refresher: &tokenRefreshTestRefresher{}},
{platform: PlatformGemini, refresher: &tokenRefreshTestRefresher{}},
{platform: PlatformAntigravity, refresher: &tokenRefreshTestRefresher{}},
},
refreshPolicy: DefaultBackgroundRefreshPolicy(),
cfg: &config.TokenRefreshConfig{RefreshBeforeExpiryHours: 1, MaxRetries: 1},
}
@@ -156,7 +181,7 @@ func TestTokenRefreshService_ProcessRefreshUsesOAuthRefreshCandidates(t *testing
svc.processRefresh()
require.Zero(t, repo.listActiveCalls, "TokenRefreshService should not use the broad active-account query")
require.Equal(t, []int64{1, 6}, repo.updatedCredentialIDs)
require.ElementsMatch(t, []int64{1, 6}, repo.updatedCredentialIDs)
require.Equal(t, 1, repo.clearTempCalls, "successful refresh should clear the OAuth 401 temp-unschedulable state")
}
@@ -5,6 +5,7 @@ package service
import (
"context"
"errors"
"reflect"
"testing"
"time"
@@ -14,21 +15,37 @@ import (
type tokenRefreshAccountRepo struct {
mockAccountRepoForGemini
updateCalls int
fullUpdateCalls int
updateCredentialsCalls int
setErrorCalls int
clearTempCalls int
setTempUnschedCalls int
updateExtraCalls int
lastErrorMessage string
lastTempUnschedReason string
lastExtraUpdates map[string]any
lastAccount *Account
updateErr error
setErrorErr error
setTempUnschedErr error
beforeConditionalState func()
updateCalls int
fullUpdateCalls int
updateCredentialsCalls int
setErrorCalls int
clearTempCalls int
setTempUnschedCalls int
updateExtraCalls int
lastErrorMessage string
lastTempUnschedReason string
lastExtraUpdates map[string]any
lastAccount *Account
updateErr error
cancelOnUpdate context.CancelFunc
conditionalErrorCalls int
conditionalTempCalls int
conditionalSuccessCalls int
conditionalErrorErr error
conditionalTempErr error
conditionalSuccessErr error
snapshotReads bool
respectReadContext bool
getByIDCalls int
durableReadDelay time.Duration
mutateSchedulingOnSuccessCAS bool
reauthorizeOnErrorCAS bool
reauthorizeOnTempCAS bool
repairProxyOnErrorCAS bool
repairProxyOnTempCAS bool
setErrorErr error
setTempUnschedErr error
beforeConditionalState func()
}
func (r *tokenRefreshAccountRepo) Update(ctx context.Context, account *Account) error {
@@ -49,13 +66,40 @@ func (r *tokenRefreshAccountRepo) UpdateCredentials(ctx context.Context, id int6
if acc, ok := r.accountsByID[id]; ok && acc != nil {
acc.Credentials = cloned
r.lastAccount = acc
if r.cancelOnUpdate != nil {
r.cancelOnUpdate()
}
return nil
}
}
r.lastAccount = &Account{ID: id, Credentials: cloned}
if r.cancelOnUpdate != nil {
r.cancelOnUpdate()
}
return nil
}
func (r *tokenRefreshAccountRepo) GetByID(ctx context.Context, id int64) (*Account, error) {
if r.respectReadContext && ctx.Err() != nil {
return nil, ctx.Err()
}
r.getByIDCalls++
if r.getByIDCalls > 1 && r.durableReadDelay > 0 {
timer := time.NewTimer(r.durableReadDelay)
defer timer.Stop()
select {
case <-ctx.Done():
return nil, ctx.Err()
case <-timer.C:
}
}
account, err := r.mockAccountRepoForGemini.GetByID(ctx, id)
if err != nil || !r.snapshotReads {
return account, err
}
return snapshotOAuthRefreshAccount(account), nil
}
func (r *tokenRefreshAccountRepo) SetError(ctx context.Context, id int64, errorMsg string) error {
r.setErrorCalls++
r.lastErrorMessage = errorMsg
@@ -132,6 +176,124 @@ func grokCredentialSnapshotMatchesAccount(account *Account, snapshot GrokCredent
grokCredentialProxyIDsEqual(account.ProxyID, snapshot.ProxyID)
}
func (r *tokenRefreshAccountRepo) SetGrokOAuthRefreshErrorIfCredentialsUnchanged(
_ context.Context,
id int64,
expectedCredentials map[string]any,
expectedProxyID *int64,
errorMsg string,
) (bool, error) {
r.conditionalErrorCalls++
if r.conditionalErrorErr != nil {
return false, r.conditionalErrorErr
}
account := r.accountsByID[id]
if account == nil {
return false, nil
}
if r.reauthorizeOnErrorCAS {
r.reauthorizeOnErrorCAS = false
account.Credentials = map[string]any{
"access_token": "fresh-access",
"refresh_token": "fresh-refresh",
"_token_version": int64(2),
}
account.Status = StatusActive
account.Schedulable = true
}
if r.repairProxyOnErrorCAS {
r.repairProxyOnErrorCAS = false
proxyID := int64(902)
account.ProxyID = &proxyID
}
if account.Status != StatusActive || account.Platform != PlatformGrok || account.Type != AccountTypeOAuth ||
!reflect.DeepEqual(account.Credentials, expectedCredentials) || !reflect.DeepEqual(account.ProxyID, expectedProxyID) {
return false, nil
}
r.setErrorCalls++
r.lastErrorMessage = errorMsg
account.Status = StatusError
account.Schedulable = false
account.ErrorMessage = errorMsg
return true, nil
}
func (r *tokenRefreshAccountRepo) UpdateGrokOAuthCredentialsIfUnchanged(
_ context.Context,
id int64,
expectedCredentials map[string]any,
expectedProxyID *int64,
credentials map[string]any,
) (bool, error) {
r.conditionalSuccessCalls++
if r.conditionalSuccessErr != nil {
return false, r.conditionalSuccessErr
}
account := r.accountsByID[id]
if account != nil && r.mutateSchedulingOnSuccessCAS {
r.mutateSchedulingOnSuccessCAS = false
account.Status = StatusDisabled
account.Schedulable = false
resetAt := time.Now().Add(30 * time.Minute)
account.RateLimitResetAt = &resetAt
}
if account == nil || account.Platform != PlatformGrok ||
account.Type != AccountTypeOAuth || !reflect.DeepEqual(account.Credentials, expectedCredentials) ||
!reflect.DeepEqual(account.ProxyID, expectedProxyID) {
return false, nil
}
r.updateCalls++
r.updateCredentialsCalls++
account.Credentials = shallowCopyMap(credentials)
r.lastAccount = account
if r.cancelOnUpdate != nil {
r.cancelOnUpdate()
}
return true, nil
}
func (r *tokenRefreshAccountRepo) SetGrokOAuthRefreshTempUnschedulableIfCredentialsUnchanged(
_ context.Context,
id int64,
expectedCredentials map[string]any,
expectedProxyID *int64,
until time.Time,
reason string,
) (bool, error) {
r.conditionalTempCalls++
if r.conditionalTempErr != nil {
return false, r.conditionalTempErr
}
account := r.accountsByID[id]
if account == nil {
return false, nil
}
if r.reauthorizeOnTempCAS {
r.reauthorizeOnTempCAS = false
account.Credentials = map[string]any{
"access_token": "fresh-access",
"refresh_token": "fresh-refresh",
"_token_version": int64(2),
}
account.Status = StatusActive
account.Schedulable = true
}
if r.repairProxyOnTempCAS {
r.repairProxyOnTempCAS = false
proxyID := int64(902)
account.ProxyID = &proxyID
}
if account.Status != StatusActive || account.Platform != PlatformGrok || account.Type != AccountTypeOAuth ||
!reflect.DeepEqual(account.Credentials, expectedCredentials) || !reflect.DeepEqual(account.ProxyID, expectedProxyID) {
return false, nil
}
r.setTempUnschedCalls++
r.lastTempUnschedReason = reason
account.TempUnschedulableUntil = &until
account.TempUnschedulableReason = reason
return true, nil
}
func (r *tokenRefreshAccountRepo) UpdateExtra(ctx context.Context, id int64, updates map[string]any) error {
r.updateExtraCalls++
r.lastExtraUpdates = shallowCopyMap(updates)
@@ -149,15 +311,46 @@ func (r *tokenRefreshAccountRepo) UpdateExtra(ctx context.Context, id int64, upd
}
type tokenCacheInvalidatorStub struct {
calls int
err error
calls int
err error
ctxErr error
lastAccount *Account
}
type tokenRefreshRuntimeBlocker struct {
blockCalls int
clearCalls int
}
func (b *tokenRefreshRuntimeBlocker) BlockAccountScheduling(*Account, time.Time, string) {
b.blockCalls++
}
func (b *tokenRefreshRuntimeBlocker) ClearAccountSchedulingBlock(int64) {
b.clearCalls++
}
func (s *tokenCacheInvalidatorStub) InvalidateToken(ctx context.Context, account *Account) error {
s.calls++
s.ctxErr = ctx.Err()
s.lastAccount = snapshotOAuthRefreshAccount(account)
return s.err
}
type tokenRefreshSchedulerCache struct {
SchedulerCache
setAccountCalls int
ctxErr error
lastAccount *Account
}
func (s *tokenRefreshSchedulerCache) SetAccount(ctx context.Context, account *Account) error {
s.setAccountCalls++
s.ctxErr = ctx.Err()
s.lastAccount = snapshotOAuthRefreshAccount(account)
return nil
}
type tempUnschedCacheStub struct {
deleteCalls int
setCalls int
@@ -182,6 +375,7 @@ func (s *tempUnschedCacheStub) DeleteTempUnsched(ctx context.Context, accountID
type tokenRefresherStub struct {
credentials map[string]any
err error
calls int
}
func (r *tokenRefresherStub) CanRefresh(account *Account) bool {
@@ -193,6 +387,7 @@ func (r *tokenRefresherStub) NeedsRefresh(account *Account, refreshWindowDuratio
}
func (r *tokenRefresherStub) Refresh(ctx context.Context, account *Account) (map[string]any, error) {
r.calls++
if r.err != nil {
return nil, r.err
}
@@ -624,7 +819,7 @@ func TestTokenRefreshService_RefreshWithRetry_AntigravityNonRetryableError(t *te
err := service.refreshWithRetry(context.Background(), account, refresher, refresher, time.Hour)
require.Error(t, err)
require.Equal(t, 0, repo.updateCalls)
require.Equal(t, 0, invalidator.calls)
require.Equal(t, 1, invalidator.calls)
require.Equal(t, 1, repo.setErrorCalls) // 不可重试错误应设置错误状态
}
@@ -763,6 +958,8 @@ type mockTokenCacheForRefreshAPI struct {
lockResult bool
lockErr error
releaseCalls int
deleteCalls int
deleteCtxErr error
}
func (m *mockTokenCacheForRefreshAPI) GetAccessToken(_ context.Context, _ string) (string, error) {
@@ -773,7 +970,9 @@ func (m *mockTokenCacheForRefreshAPI) SetAccessToken(_ context.Context, _ string
return nil
}
func (m *mockTokenCacheForRefreshAPI) DeleteAccessToken(_ context.Context, _ string) error {
func (m *mockTokenCacheForRefreshAPI) DeleteAccessToken(ctx context.Context, _ string) error {
m.deleteCalls++
m.deleteCtxErr = ctx.Err()
return nil
}
@@ -788,6 +987,11 @@ func (m *mockTokenCacheForRefreshAPI) ReleaseRefreshLock(_ context.Context, _ st
// buildPathAService 构建注入了 refreshAPI 的 servicePath A 测试辅助)
func buildPathAService(repo *tokenRefreshAccountRepo, cache GeminiTokenCache, invalidator TokenCacheInvalidator) (*TokenRefreshService, *tokenRefresherStub) {
for _, account := range repo.accountsByID {
if account != nil && account.Status == "" {
account.Status = StatusActive
}
}
cfg := &config.Config{
TokenRefresh: config.TokenRefreshConfig{
MaxRetries: 1,
@@ -828,6 +1032,210 @@ func TestPathA_Success(t *testing.T) {
require.Equal(t, 1, cache.releaseCalls) // 锁被释放
}
func TestPathA_GrokSuccessPersistenceFailureContainsProviderWithoutRetryOrMutation(t *testing.T) {
account := &Account{
ID: 110,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Credentials: map[string]any{
"access_token": "attempted-access",
"refresh_token": "attempted-refresh",
},
}
repo := &tokenRefreshAccountRepo{
conditionalSuccessErr: errors.New("database unavailable after provider success"),
}
repo.accountsByID = map[int64]*Account{account.ID: account}
cfg := &config.Config{TokenRefresh: config.TokenRefreshConfig{
MaxRetries: 3,
RetryBackoffSeconds: 0,
}}
svc := NewTokenRefreshService(repo, nil, nil, nil, nil, nil, nil, cfg, nil)
svc.SetRefreshAPI(NewOAuthRefreshAPI(repo, nil))
refresher := &tokenRefresherStub{credentials: map[string]any{
"access_token": "provider-access",
"refresh_token": "provider-refresh",
}}
err := svc.refreshWithRetry(context.Background(), account, refresher, refresher, time.Hour)
var containmentErr *providerCycleContainmentRefreshError
require.ErrorAs(t, err, &containmentErr)
require.Equal(t, 1, refresher.calls, "a provider-issued rotated token must never be retried after persistence fails")
require.Equal(t, 1, repo.conditionalSuccessCalls)
require.Zero(t, repo.conditionalErrorCalls)
require.Zero(t, repo.conditionalTempCalls)
require.Equal(t, StatusActive, account.Status)
require.Equal(t, "attempted-refresh", account.GetGrokRefreshToken())
}
func TestPathA_GrokSuccessPublishesDurableSchedulingState(t *testing.T) {
account := &Account{
ID: 111,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Credentials: map[string]any{
"access_token": "attempted-access",
"refresh_token": "attempted-refresh",
},
}
repo := &tokenRefreshAccountRepo{
snapshotReads: true,
mutateSchedulingOnSuccessCAS: true,
}
repo.accountsByID = map[int64]*Account{account.ID: account}
scheduler := &tokenRefreshSchedulerCache{}
cfg := &config.Config{TokenRefresh: config.TokenRefreshConfig{MaxRetries: 1}}
svc := NewTokenRefreshService(repo, nil, nil, nil, nil, nil, scheduler, cfg, nil)
svc.SetRefreshAPI(NewOAuthRefreshAPI(repo, nil))
refresher := &tokenRefresherStub{credentials: map[string]any{
"access_token": "provider-access",
"refresh_token": "provider-refresh",
}}
err := svc.refreshWithRetry(context.Background(), account, refresher, refresher, time.Hour)
require.NoError(t, err)
require.Equal(t, StatusDisabled, repo.accountsByID[account.ID].Status)
require.False(t, repo.accountsByID[account.ID].Schedulable)
require.NotNil(t, repo.accountsByID[account.ID].RateLimitResetAt)
require.Equal(t, 1, scheduler.setAccountCalls)
require.NotNil(t, scheduler.lastAccount)
require.Equal(t, StatusDisabled, scheduler.lastAccount.Status)
require.False(t, scheduler.lastAccount.Schedulable)
require.NotNil(t, scheduler.lastAccount.RateLimitResetAt,
"post-refresh cache publication must preserve the durable concurrent exclusion state")
}
func TestPathA_GrokCancelAfterSuccessCASUsesDetachedDurableStateAndInvalidatesCache(t *testing.T) {
account := &Account{
ID: 112,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Credentials: map[string]any{
"access_token": "attempted-access",
"refresh_token": "attempted-refresh",
},
}
ctx, cancel := context.WithCancel(context.Background())
repo := &tokenRefreshAccountRepo{
cancelOnUpdate: cancel,
snapshotReads: true,
respectReadContext: true,
mutateSchedulingOnSuccessCAS: true,
}
repo.accountsByID = map[int64]*Account{account.ID: account}
invalidator := &tokenCacheInvalidatorStub{}
scheduler := &tokenRefreshSchedulerCache{}
cache := &mockTokenCacheForRefreshAPI{lockResult: true}
cfg := &config.Config{TokenRefresh: config.TokenRefreshConfig{MaxRetries: 1}}
svc := NewTokenRefreshService(repo, nil, nil, nil, nil, invalidator, scheduler, cfg, nil)
svc.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache))
refresher := &tokenRefresherStub{credentials: map[string]any{
"access_token": "provider-access",
"refresh_token": "provider-refresh",
}}
err := svc.refreshWithRetry(ctx, account, refresher, refresher, time.Hour)
require.ErrorIs(t, err, context.Canceled)
require.Equal(t, 1, repo.conditionalSuccessCalls)
require.Equal(t, "provider-refresh", repo.accountsByID[account.ID].GetGrokRefreshToken())
require.Equal(t, 1, cache.deleteCalls)
require.NoError(t, cache.deleteCtxErr)
require.Equal(t, 1, invalidator.calls, "the pre-rotation access-token cache must be invalidated after committed CAS")
require.NoError(t, invalidator.ctxErr)
require.NotNil(t, invalidator.lastAccount)
require.Equal(t, "provider-refresh", invalidator.lastAccount.GetGrokRefreshToken())
require.Equal(t, StatusDisabled, invalidator.lastAccount.Status)
require.Equal(t, 1, scheduler.setAccountCalls)
require.NoError(t, scheduler.ctxErr)
require.NotNil(t, scheduler.lastAccount)
require.Equal(t, StatusDisabled, scheduler.lastAccount.Status)
require.False(t, scheduler.lastAccount.Schedulable)
require.NotNil(t, scheduler.lastAccount.RateLimitResetAt)
}
func TestTokenRefreshService_PersistedSuccessCrossingAttemptDeadlineStaysSuccessful(t *testing.T) {
account := &Account{
ID: 113,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Credentials: map[string]any{
"access_token": "attempted-access",
"refresh_token": "attempted-refresh",
},
}
repo := &tokenRefreshAccountRepo{
snapshotReads: true,
durableReadDelay: 30 * time.Millisecond,
}
repo.accountsByID = map[int64]*Account{account.ID: account}
scheduler := &tokenRefreshSchedulerCache{}
svc := &TokenRefreshService{
accountRepo: repo,
refreshAPI: NewOAuthRefreshAPI(repo, nil),
refreshPolicy: DefaultBackgroundRefreshPolicy(),
cfg: &config.TokenRefreshConfig{MaxRetries: 1, ProviderFailureThreshold: 1},
schedulerCache: scheduler,
attemptTimeoutOverride: 10 * time.Millisecond,
}
refresher := &tokenRefresherStub{credentials: map[string]any{
"access_token": "provider-access",
"refresh_token": "provider-refresh",
}}
state := &tokenRefreshProviderState{
service: svc,
rateGate: newTokenRefreshRateGate(10000),
poolGate: newTokenRefreshConcurrencyGate(1),
}
err := svc.refreshWithRetryWithRateGate(context.Background(), account, refresher, refresher, time.Hour, state)
state.recordResult(err)
require.NoError(t, err)
require.Equal(t, 1, refresher.calls, "durably persisted success must not retry after only the internal attempt deadline elapsed")
require.Equal(t, 1, repo.conditionalSuccessCalls)
require.Zero(t, repo.conditionalTempCalls)
require.Zero(t, repo.setTempUnschedCalls)
require.False(t, state.isTripped(), "a durable success must not count toward the provider breaker")
require.Equal(t, "provider-refresh", repo.accountsByID[account.ID].GetGrokRefreshToken())
require.Equal(t, 1, scheduler.setAccountCalls)
}
func TestPathA_ParentCancellationAfterPersistStillSynchronizesCacheState(t *testing.T) {
account := &Account{
ID: 109,
Platform: PlatformGemini,
Type: AccountTypeOAuth,
Status: StatusActive,
}
ctx, cancel := context.WithCancel(context.Background())
repo := &tokenRefreshAccountRepo{cancelOnUpdate: cancel}
repo.accountsByID = map[int64]*Account{account.ID: account}
invalidator := &tokenCacheInvalidatorStub{}
scheduler := &tokenRefreshSchedulerCache{}
cache := &mockTokenCacheForRefreshAPI{lockResult: true}
service, refresher := buildPathAService(repo, cache, invalidator)
service.schedulerCache = scheduler
err := service.refreshWithRetry(ctx, account, refresher, refresher, time.Hour)
require.ErrorIs(t, err, context.Canceled)
require.Equal(t, 1, repo.updateCredentialsCalls, "credentials were durably persisted before cancellation")
require.Equal(t, 1, invalidator.calls)
require.NoError(t, invalidator.ctxErr, "post-persist invalidation must use bounded cleanup context")
require.Equal(t, 1, scheduler.setAccountCalls)
require.NoError(t, scheduler.ctxErr, "scheduler sync must use bounded cleanup context")
}
// TestPathA_LockHeld 锁被其他 worker 持有 → 返回 errRefreshSkipped
func TestPathA_LockHeld(t *testing.T) {
account := &Account{
@@ -912,7 +1320,7 @@ func TestPathA_NonRetryableError(t *testing.T) {
require.Error(t, err)
require.Equal(t, 1, repo.setErrorCalls) // 应标记 error 状态
require.Equal(t, 0, repo.updateCalls) // 不应更新 credentials
require.Equal(t, 0, invalidator.calls) // 不应触发缓存失效
require.Equal(t, 1, invalidator.calls) // 永久凭证失败后必须失效旧 token 缓存
}
// TestPathA_RetryableErrorExhausted 统一 API 路径可重试错误耗尽 → 不标记 error
@@ -949,6 +1357,256 @@ func TestPathA_RetryableErrorExhausted(t *testing.T) {
require.Equal(t, 0, invalidator.calls) // 不应触发缓存失效
}
func TestPathA_GrokPermanentFailureCASLetsConcurrentAccountRepairWin(t *testing.T) {
tests := []struct {
name string
configure func(*tokenRefreshAccountRepo)
assert func(*testing.T, *Account)
}{
{
name: "credential reauthorization",
configure: func(repo *tokenRefreshAccountRepo) {
repo.reauthorizeOnErrorCAS = true
},
assert: func(t *testing.T, account *Account) {
require.Equal(t, "fresh-refresh", account.GetGrokRefreshToken())
},
},
{
name: "proxy repair",
configure: func(repo *tokenRefreshAccountRepo) {
repo.repairProxyOnErrorCAS = true
},
assert: func(t *testing.T, account *Account) {
require.NotNil(t, account.ProxyID)
require.Equal(t, int64(902), *account.ProxyID)
require.Equal(t, "attempted-refresh", account.GetGrokRefreshToken(),
"proxy-only repair must prove the proxy fingerprint independently of credentials")
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
proxyID := int64(901)
account := &Account{
ID: 120,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
ProxyID: &proxyID,
Credentials: map[string]any{
"access_token": "attempted-access",
"refresh_token": "attempted-refresh",
"_token_version": int64(1),
},
}
repo := &tokenRefreshAccountRepo{}
repo.accountsByID = map[int64]*Account{account.ID: account}
tt.configure(repo)
invalidator := &tokenCacheInvalidatorStub{}
cache := &mockTokenCacheForRefreshAPI{lockResult: true}
service, _ := buildPathAService(repo, cache, invalidator)
blocker := &tokenRefreshRuntimeBlocker{}
service.SetAccountRuntimeBlocker(blocker)
refresher := &tokenRefresherStub{err: errors.New("invalid_grant: revoked")}
err := service.refreshWithRetry(context.Background(), account, refresher, refresher, time.Hour)
require.ErrorIs(t, err, errRefreshSkipped)
require.Equal(t, 1, repo.conditionalErrorCalls)
require.Zero(t, repo.setErrorCalls)
require.Zero(t, blocker.blockCalls)
require.Zero(t, invalidator.calls, "a stale permanent failure must not invalidate newly repaired credentials")
require.Equal(t, StatusActive, account.Status)
require.True(t, account.Schedulable)
tt.assert(t, account)
})
}
}
func TestPathA_GrokTransientFailureCASLetsConcurrentAccountRepairWin(t *testing.T) {
tests := []struct {
name string
configure func(*tokenRefreshAccountRepo)
assert func(*testing.T, *Account)
}{
{
name: "credential reauthorization",
configure: func(repo *tokenRefreshAccountRepo) {
repo.reauthorizeOnTempCAS = true
},
assert: func(t *testing.T, account *Account) {
require.Equal(t, "fresh-refresh", account.GetGrokRefreshToken())
},
},
{
name: "proxy repair",
configure: func(repo *tokenRefreshAccountRepo) {
repo.repairProxyOnTempCAS = true
},
assert: func(t *testing.T, account *Account) {
require.NotNil(t, account.ProxyID)
require.Equal(t, int64(902), *account.ProxyID)
require.Equal(t, "attempted-refresh", account.GetGrokRefreshToken(),
"proxy-only repair must prove the proxy fingerprint independently of credentials")
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
proxyID := int64(901)
account := &Account{
ID: 121,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
ProxyID: &proxyID,
Credentials: map[string]any{
"access_token": "attempted-access",
"refresh_token": "attempted-refresh",
"_token_version": int64(1),
},
}
repo := &tokenRefreshAccountRepo{}
repo.accountsByID = map[int64]*Account{account.ID: account}
tt.configure(repo)
invalidator := &tokenCacheInvalidatorStub{}
cache := &mockTokenCacheForRefreshAPI{lockResult: true}
service, _ := buildPathAService(repo, cache, invalidator)
blocker := &tokenRefreshRuntimeBlocker{}
service.SetAccountRuntimeBlocker(blocker)
refresher := &tokenRefresherStub{err: errors.New("temporary provider timeout")}
err := service.refreshWithRetry(context.Background(), account, refresher, refresher, time.Hour)
require.ErrorIs(t, err, errRefreshSkipped)
require.Equal(t, 1, repo.conditionalTempCalls)
require.Zero(t, repo.setTempUnschedCalls)
require.Zero(t, blocker.blockCalls)
require.Equal(t, StatusActive, account.Status)
require.True(t, account.Schedulable)
require.Nil(t, account.TempUnschedulableUntil)
tt.assert(t, account)
})
}
}
func TestTokenRefreshService_GrokMissingConditionalMutationContractContainsProviderCycle(t *testing.T) {
tests := []struct {
name string
refreshErr error
}{
{name: "permanent failure", refreshErr: errors.New("invalid_grant: revoked")},
{name: "transient failure", refreshErr: errors.New("temporary provider timeout")},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
svc := &TokenRefreshService{
accountRepo: &mockAccountRepoForGemini{},
refreshPolicy: DefaultBackgroundRefreshPolicy(),
cfg: &config.TokenRefreshConfig{MaxRetries: 1},
}
account := &Account{
ID: 122,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Credentials: map[string]any{"refresh_token": "attempted"},
}
refresher := &tokenRefresherStub{err: tt.refreshErr}
err := svc.refreshWithRetry(context.Background(), account, refresher, nil, time.Hour)
var providerErr *providerConfigurationRefreshError
require.ErrorAs(t, err, &providerErr)
state := &tokenRefreshProviderState{service: svc}
state.recordResult(err)
require.True(t, state.isTripped(), "a missing safety contract must stop the provider cycle")
require.Equal(t, StatusActive, account.Status)
require.True(t, account.Schedulable)
})
}
}
func TestTokenRefreshService_GrokConditionalMutationErrorsContainProviderCycle(t *testing.T) {
tests := []struct {
name string
upstreamErr error
configureRepo func(*tokenRefreshAccountRepo, error)
expectedCASCalls func(*tokenRefreshAccountRepo) int
}{
{
name: "permanent failure",
upstreamErr: errors.New("invalid_grant: revoked"),
configureRepo: func(repo *tokenRefreshAccountRepo, casErr error) {
repo.conditionalErrorErr = casErr
},
expectedCASCalls: func(repo *tokenRefreshAccountRepo) int { return repo.conditionalErrorCalls },
},
{
name: "transient failure",
upstreamErr: errors.New("temporary provider timeout"),
configureRepo: func(repo *tokenRefreshAccountRepo, casErr error) {
repo.conditionalTempErr = casErr
},
expectedCASCalls: func(repo *tokenRefreshAccountRepo) int { return repo.conditionalTempCalls },
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
account := &Account{
ID: 123,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Credentials: map[string]any{"refresh_token": "attempted"},
}
casErr := errors.New("conditional account mutation unavailable")
repo := &tokenRefreshAccountRepo{}
repo.accountsByID = map[int64]*Account{account.ID: account}
tt.configureRepo(repo, casErr)
invalidator := &tokenCacheInvalidatorStub{}
blocker := &tokenRefreshRuntimeBlocker{}
svc := &TokenRefreshService{
accountRepo: repo,
refreshPolicy: DefaultBackgroundRefreshPolicy(),
cfg: &config.TokenRefreshConfig{MaxRetries: 1},
cacheInvalidator: invalidator,
}
svc.SetAccountRuntimeBlocker(blocker)
refresher := &tokenRefresherStub{err: tt.upstreamErr}
err := svc.refreshWithRetry(context.Background(), account, refresher, nil, time.Hour)
var containmentErr *providerCycleContainmentRefreshError
require.ErrorAs(t, err, &containmentErr)
require.ErrorIs(t, err, casErr)
require.NotErrorIs(t, err, tt.upstreamErr, "a CAS execution failure must replace the stale upstream classification")
var permanentErr *accountPermanentRefreshError
require.False(t, errors.As(err, &permanentErr))
require.Equal(t, 1, tt.expectedCASCalls(repo))
state := &tokenRefreshProviderState{service: svc}
state.recordResult(err)
require.True(t, state.isTripped(), "an unsafe mutation result must stop the provider cycle immediately")
require.Zero(t, repo.setErrorCalls)
require.Zero(t, repo.setTempUnschedCalls)
require.Zero(t, blocker.blockCalls)
require.Zero(t, invalidator.calls)
require.Equal(t, StatusActive, account.Status)
require.True(t, account.Schedulable)
})
}
}
// TestPathA_DBUpdateFailed 统一 API 路径 DB 更新失败 → 返回 error,不执行 postRefreshActions
func TestPathA_DBUpdateFailed(t *testing.T) {
account := &Account{
+1
View File
@@ -687,6 +687,7 @@ var ProviderSet = wire.NewSet(
NewCRSSyncService,
ProvideUpdateService,
ProvideTokenRefreshService,
wire.Bind(new(GrokOAuthReconciler), new(*TokenRefreshService)),
ProvideAccountExpiryService,
ProvideProxyExpiryService,
ProvideSubscriptionExpiryService,
+18
View File
@@ -664,6 +664,24 @@ token_refresh:
# Whether OpenAI refresh flow is allowed to sync linked Sora accounts
# 是否允许 OpenAI 刷新流程同步覆盖 linked_openai_account_id 关联的 Sora 账号 token
sync_linked_sora_accounts: false
# Candidate accounts loaded per cursor page (maximum 1000)
# 每个游标分页加载的候选账号数量(最大 1000)
candidate_page_size: 200
# Maximum concurrent refresh attempts per provider (maximum 32)
# 每个平台的最大并发刷新数(最大 32)
provider_concurrency: 4
# Per-provider refresh requests per second in each server process (maximum 100)
# 每个服务进程中每个平台每秒允许的刷新请求数(最大 100)
provider_qps: 2
# Consecutive transient failures that contain a provider for the current cycle (maximum 100)
# 当前周期内触发平台级熔断的连续临时失败次数(最大 100)
provider_failure_threshold: 3
# Timeout for one upstream refresh attempt, in seconds (maximum 300)
# 单次上游刷新尝试的超时时间(秒,最大 300)
attempt_timeout_seconds: 15
# Total timeout for one background refresh cycle, in seconds (maximum 3600)
# 单个后台刷新周期的总超时时间(秒,最大 3600)
cycle_timeout_seconds: 240
# =============================================================================
# API Key Auth Cache Configuration