diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index 4012e5d228..1e28c09153 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -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) diff --git a/backend/internal/config/config.go b/backend/internal/config/config.go index 40947ef952..2064c0c7ea 100644 --- a/backend/internal/config/config.go +++ b/backend/internal/config/config.go @@ -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 diff --git a/backend/internal/handler/admin/grok_import_probe_handler_test.go b/backend/internal/handler/admin/grok_import_probe_handler_test.go index 489a13ae6d..bc8c3f3f6a 100644 --- a/backend/internal/handler/admin/grok_import_probe_handler_test.go +++ b/backend/internal/handler/admin/grok_import_probe_handler_test.go @@ -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() diff --git a/backend/internal/handler/admin/grok_oauth_handler.go b/backend/internal/handler/admin/grok_oauth_handler.go index 1a309b7c9a..258c68853d 100644 --- a/backend/internal/handler/admin/grok_oauth_handler.go +++ b/backend/internal/handler/admin/grok_oauth_handler.go @@ -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"` diff --git a/backend/internal/handler/admin/grok_oauth_handler_test.go b/backend/internal/handler/admin/grok_oauth_handler_test.go index 9065e3321b..ee7e8e5e3e 100644 --- a/backend/internal/handler/admin/grok_oauth_handler_test.go +++ b/backend/internal/handler/admin/grok_oauth_handler_test.go @@ -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) +} diff --git a/backend/internal/repository/account_repo.go b/backend/internal/repository/account_repo.go index 9af99223f7..04beff8c62 100644 --- a/backend/internal/repository/account_repo.go +++ b/backend/internal/repository/account_repo.go @@ -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 在账号状态变更时主动同步快照到调度器缓存。 // 当账号被设置为错误、禁用、不可调度或临时不可调度时调用， // 确保调度器和粘性会话逻辑能及时感知账号的最新状态，避免继续使用不可用账号。 diff --git a/backend/internal/repository/account_repo_integration_test.go b/backend/internal/repository/account_repo_integration_test.go index b59ce2077f..fa4eed7f5c 100644 --- a/backend/internal/repository/account_repo_integration_test.go +++ b/backend/internal/repository/account_repo_integration_test.go @@ -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() { diff --git a/backend/internal/repository/account_repo_temp_unsched_test.go b/backend/internal/repository/account_repo_temp_unsched_test.go index 2ea18bd214..df1e7abf3b 100644 --- a/backend/internal/repository/account_repo_temp_unsched_test.go +++ b/backend/internal/repository/account_repo_temp_unsched_test.go @@ -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...) } diff --git a/backend/internal/server/routes/admin.go b/backend/internal/server/routes/admin.go index 6bf67bcbb1..a60adb1826 100644 --- a/backend/internal/server/routes/admin.go +++ b/backend/internal/server/routes/admin.go @@ -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) diff --git a/backend/internal/service/account_service.go b/backend/internal/service/account_service.go index 109ef95680..b7fe0a68cc 100644 --- a/backend/internal/service/account_service.go +++ b/backend/internal/service/account_service.go @@ -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 diff --git a/backend/internal/service/grok_credential_failure.go b/backend/internal/service/grok_credential_failure.go index 65dd844072..4ac4aec084 100644 --- a/backend/internal/service/grok_credential_failure.go +++ b/backend/internal/service/grok_credential_failure.go @@ -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"), diff --git a/backend/internal/service/grok_credential_failure_test.go b/backend/internal/service/grok_credential_failure_test.go index c32559ea12..b493092927 100644 --- a/backend/internal/service/grok_credential_failure_test.go +++ b/backend/internal/service/grok_credential_failure_test.go @@ -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) diff --git a/backend/internal/service/grok_oauth_reconciliation.go b/backend/internal/service/grok_oauth_reconciliation.go new file mode 100644 index 0000000000..ca7611191f --- /dev/null +++ b/backend/internal/service/grok_oauth_reconciliation.go @@ -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 +} diff --git a/backend/internal/service/grok_oauth_reconciliation_test.go b/backend/internal/service/grok_oauth_reconciliation_test.go new file mode 100644 index 0000000000..8de04a1b5d --- /dev/null +++ b/backend/internal/service/grok_oauth_reconciliation_test.go @@ -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) +} diff --git a/backend/internal/service/grok_token_refresher.go b/backend/internal/service/grok_token_refresher.go index 0b87fcac78..d2d4d534fc 100644 --- a/backend/internal/service/grok_token_refresher.go +++ b/backend/internal/service/grok_token_refresher.go @@ -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 diff --git a/backend/internal/service/oauth_refresh_api.go b/backend/internal/service/oauth_refresh_api.go index b80cf84b64..c940bd335a 100644 --- a/backend/internal/service/oauth_refresh_api.go +++ b/backend/internal/service/oauth_refresh_api.go @@ -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 // 刷新后的 credentials（nil 表示未刷新） - 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") diff --git a/backend/internal/service/oauth_refresh_api_test.go b/backend/internal/service/oauth_refresh_api_test.go index adf089f0bd..e92bd99f02 100644 --- a/backend/internal/service/oauth_refresh_api_test.go +++ b/backend/internal/service/oauth_refresh_api_test.go @@ -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 © } 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) { diff --git a/backend/internal/service/token_refresh_pool_health_test.go b/backend/internal/service/token_refresh_pool_health_test.go new file mode 100644 index 0000000000..f4950d879f --- /dev/null +++ b/backend/internal/service/token_refresh_pool_health_test.go @@ -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: request blocked`)} + 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()) +} diff --git a/backend/internal/service/token_refresh_service.go b/backend/internal/service/token_refresh_service.go index 2d4a026692..31d2e19909 100644 --- a/backend/internal/service/token_refresh_service.go +++ b/backend/internal/service/token_refresh_service.go @@ -10,18 +10,54 @@ import ( "time" "github.com/Wei-Shaw/sub2api/internal/config" + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" "github.com/Wei-Shaw/sub2api/internal/util/logredact" ) // tokenRefreshTempUnschedDuration token 刷新重试耗尽后临时不可调度的持续时间 const tokenRefreshTempUnschedDuration = 10 * time.Minute +const ( + defaultTokenRefreshCandidatePageSize = 200 + maxTokenRefreshCandidatePageSize = 1000 + defaultTokenRefreshProviderConcurrency = 4 + maxTokenRefreshProviderConcurrency = 32 + defaultTokenRefreshProviderQPS = 2 + maxTokenRefreshProviderQPS = 100 + defaultTokenRefreshProviderFailureThreshold = 3 + maxTokenRefreshProviderFailureThreshold = 100 + defaultTokenRefreshMaxRetries = 1 + maxTokenRefreshMaxRetries = 10 + maxTokenRefreshRetryBackoff = 30 * time.Second + defaultTokenRefreshAttemptTimeout = 15 * time.Second + maxTokenRefreshAttemptTimeout = 5 * time.Minute + maxTokenRefreshLockSafetyMargin = 5 * time.Second + defaultTokenRefreshCycleTimeout = 4 * time.Minute + maxTokenRefreshCycleTimeout = time.Hour + defaultTokenRefreshCleanupTimeout = 2 * time.Second +) + +type tokenRefreshRegistration struct { + platform string + refresher TokenRefresher + executor OAuthRefreshExecutor +} + +// GrokOAuthRefreshMutationRepository protects background refresh failure +// mutations with the exact credential document used by the upstream attempt. +// This contract is intentionally Grok-only; existing provider behavior remains +// unchanged. +type GrokOAuthRefreshMutationRepository interface { + SetGrokOAuthRefreshErrorIfCredentialsUnchanged(ctx context.Context, id int64, expectedCredentials map[string]any, expectedProxyID *int64, errorMsg string) (bool, error) + SetGrokOAuthRefreshTempUnschedulableIfCredentialsUnchanged(ctx context.Context, id int64, expectedCredentials map[string]any, expectedProxyID *int64, until time.Time, reason string) (bool, error) +} + // TokenRefreshService OAuth token自动刷新服务 // 定期检查并刷新即将过期的token type TokenRefreshService struct { accountRepo AccountRepository - refreshers []TokenRefresher - executors []OAuthRefreshExecutor // 与 refreshers 一一对应的 executor（带 CacheKey） + candidatePager OAuthRefreshCandidatePager + registrations []tokenRefreshRegistration refreshPolicy BackgroundRefreshPolicy cfg *config.TokenRefreshConfig cacheInvalidator TokenCacheInvalidator @@ -34,9 +70,19 @@ type TokenRefreshService struct { privacyClientFactory PrivacyClientFactory proxyRepo ProxyRepository - stopCh chan struct{} - stopOnce sync.Once - wg sync.WaitGroup + stopCh chan struct{} + stopOnce sync.Once + wg sync.WaitGroup + runCtx context.Context + runCancel context.CancelFunc + candidateMu sync.Mutex + afterID int64 + providerMu sync.Mutex + providerGates map[string]*tokenRefreshRateGate + providerPools map[string]*tokenRefreshConcurrencyGate + + // Test-only duration seam; production uses TokenRefreshConfig seconds. + attemptTimeoutOverride time.Duration } // NewTokenRefreshService 创建token刷新服务 @@ -52,14 +98,24 @@ func NewTokenRefreshService( tempUnschedCache TempUnschedCache, grokOAuthServices ...*GrokOAuthService, ) *TokenRefreshService { + refreshCfg := &config.TokenRefreshConfig{} + if cfg != nil { + refreshCfg = &cfg.TokenRefresh + } + runCtx, runCancel := context.WithCancel(context.Background()) s := &TokenRefreshService{ accountRepo: accountRepo, refreshPolicy: DefaultBackgroundRefreshPolicy(), - cfg: &cfg.TokenRefresh, + cfg: refreshCfg, cacheInvalidator: cacheInvalidator, schedulerCache: schedulerCache, tempUnschedCache: tempUnschedCache, stopCh: make(chan struct{}), + runCtx: runCtx, + runCancel: runCancel, + } + if pager, ok := accountRepo.(OAuthRefreshCandidatePager); ok { + s.candidatePager = pager } openAIRefresher := NewOpenAITokenRefresher(openaiOAuthService, accountRepo) @@ -73,27 +129,41 @@ func NewTokenRefreshService( } grokRefresher := NewGrokTokenRefresher(grokOAuthService) - // 注册平台特定的刷新器（TokenRefresher 接口） - s.refreshers = []TokenRefresher{ - claudeRefresher, - openAIRefresher, - geminiRefresher, - agRefresher, - grokRefresher, - } - - // 注册对应的 OAuthRefreshExecutor（带 CacheKey 方法） - s.executors = []OAuthRefreshExecutor{ - claudeRefresher, - openAIRefresher, - geminiRefresher, - agRefresher, - grokRefresher, + // Each provider is registered exactly once. The same registry supplies both + // execution and repository eligibility, preventing future platform drift. + s.registrations = []tokenRefreshRegistration{ + {platform: PlatformAnthropic, refresher: claudeRefresher, executor: claudeRefresher}, + {platform: PlatformOpenAI, refresher: openAIRefresher, executor: openAIRefresher}, + {platform: PlatformGemini, refresher: geminiRefresher, executor: geminiRefresher}, + {platform: PlatformAntigravity, refresher: agRefresher, executor: agRefresher}, + {platform: PlatformGrok, refresher: grokRefresher, executor: grokRefresher}, } return s } +func (s *TokenRefreshService) eligiblePlatforms() []string { + platforms := make([]string, 0, len(s.registrations)) + for _, registration := range s.registrations { + if registration.platform != "" && registration.refresher != nil { + platforms = append(platforms, registration.platform) + } + } + return platforms +} + +func (s *TokenRefreshService) candidateAfterID() int64 { + s.candidateMu.Lock() + defer s.candidateMu.Unlock() + return s.afterID +} + +func (s *TokenRefreshService) setCandidateAfterID(afterID int64) { + s.candidateMu.Lock() + s.afterID = afterID + s.candidateMu.Unlock() +} + // SetPrivacyDeps 注入 OpenAI privacy opt-out 所需依赖 func (s *TokenRefreshService) SetPrivacyDeps(factory PrivacyClientFactory, proxyRepo ProxyRepository) { s.privacyClientFactory = factory @@ -130,7 +200,7 @@ func (s *TokenRefreshService) notifyAccountSchedulingBlockCleared(accountID int6 // Start 启动后台刷新服务 func (s *TokenRefreshService) Start() { - if !s.cfg.Enabled { + if s.cfg == nil || !s.cfg.Enabled { slog.Info("token_refresh.service_disabled") return } @@ -147,6 +217,9 @@ func (s *TokenRefreshService) Start() { // Stop 停止刷新服务（可安全多次调用） func (s *TokenRefreshService) Stop() { s.stopOnce.Do(func() { + if s.runCancel != nil { + s.runCancel() + } close(s.stopCh) }) s.wg.Wait() @@ -156,6 +229,10 @@ func (s *TokenRefreshService) Stop() { // refreshLoop 刷新循环 func (s *TokenRefreshService) refreshLoop() { defer s.wg.Done() + ctx := s.runCtx + if ctx == nil { + ctx = context.Background() + } // 计算检查间隔 checkInterval := time.Duration(s.cfg.CheckIntervalMinutes) * time.Minute @@ -167,183 +244,841 @@ func (s *TokenRefreshService) refreshLoop() { defer ticker.Stop() // 启动时立即执行一次检查 - s.processRefresh() + s.processRefreshContext(ctx) for { select { case <-ticker.C: - s.processRefresh() + s.processRefreshContext(ctx) + case <-ctx.Done(): + return case <-s.stopCh: return } } } -// processRefresh 执行一次刷新检查 -func (s *TokenRefreshService) processRefresh() { - ctx := context.Background() +type tokenRefreshPageStats struct { + total int + oauth int + needsRefresh int + refreshed int + skipped int + failed int +} - // 计算刷新窗口 - refreshWindow := time.Duration(s.cfg.RefreshBeforeExpiryHours * float64(time.Hour)) +type tokenRefreshProviderState struct { + service *TokenRefreshService + registration tokenRefreshRegistration + rateGate refreshAttemptGate + poolGate *tokenRefreshConcurrencyGate - // 获取所有active状态的账号 - accounts, err := s.listActiveAccounts(ctx) + mu sync.Mutex + consecutiveFailures int + tripped bool +} + +type tokenRefreshRateGate struct { + mu sync.Mutex + next time.Time + interval time.Duration +} + +type tokenRefreshConcurrencyGate struct { + slots chan struct{} +} + +type refreshAttemptGate interface { + acquire(ctx context.Context) (release func(), err error) +} + +type providerRefreshAttemptGate interface { + refreshAttemptGate + acquireRate(ctx context.Context) (release func(), err error) +} + +type rateLimitedOAuthRefreshExecutor struct { + OAuthRefreshExecutor + acquireRate func(context.Context) (func(), error) +} + +func (e *rateLimitedOAuthRefreshExecutor) Refresh(ctx context.Context, account *Account) (map[string]any, error) { + if e == nil || e.OAuthRefreshExecutor == nil { + return nil, errors.New("OAuth refresh executor is not configured") + } + release := func() {} + if e.acquireRate != nil { + var err error + release, err = e.acquireRate(ctx) + if err != nil { + return nil, err + } + } + defer release() + return e.OAuthRefreshExecutor.Refresh(ctx, account) +} + +func newTokenRefreshRateGate(qps int) *tokenRefreshRateGate { + if qps <= 0 { + return &tokenRefreshRateGate{} + } + return newTokenRefreshRateGateWithInterval(time.Second / time.Duration(qps)) +} + +// newTokenRefreshRateGateWithInterval is a narrow duration seam used to test +// slot reservation and cancellation without waiting on production-scale QPS. +func newTokenRefreshRateGateWithInterval(interval time.Duration) *tokenRefreshRateGate { + return &tokenRefreshRateGate{interval: interval} +} + +func (g *tokenRefreshRateGate) reserveSlot(now time.Time) time.Time { + g.mu.Lock() + defer g.mu.Unlock() + if g.next.Before(now) { + g.next = now + } + slot := g.next + g.next = g.next.Add(g.interval) + return slot +} + +func (g *tokenRefreshRateGate) wait(ctx context.Context) error { + if g == nil || g.interval <= 0 { + return nil + } + slot := g.reserveSlot(time.Now()) + + wait := time.Until(slot) + if wait <= 0 { + return nil + } + timer := time.NewTimer(wait) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + +func (g *tokenRefreshRateGate) acquire(ctx context.Context) (func(), error) { + if err := g.wait(ctx); err != nil { + return nil, err + } + return func() {}, nil +} + +func newTokenRefreshConcurrencyGate(concurrency int) *tokenRefreshConcurrencyGate { + if concurrency < 1 { + concurrency = 1 + } + return &tokenRefreshConcurrencyGate{slots: make(chan struct{}, concurrency)} +} + +func (g *tokenRefreshConcurrencyGate) acquire(ctx context.Context) (func(), error) { + if g == nil { + return func() {}, nil + } + select { + case g.slots <- struct{}{}: + return func() { <-g.slots }, nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +func (p *tokenRefreshProviderState) isTripped() bool { + p.mu.Lock() + defer p.mu.Unlock() + return p.tripped +} + +func (p *tokenRefreshProviderState) acquire(ctx context.Context) (func(), error) { + if p == nil || p.isTripped() { + return nil, errRefreshSkipped + } + release, err := p.poolGate.acquire(ctx) if err != nil { - slog.Error("token_refresh.list_accounts_failed", "error", err) + return nil, err + } + if p.isTripped() { + release() + return nil, errRefreshSkipped + } + return release, nil +} + +func (p *tokenRefreshProviderState) acquireRate(ctx context.Context) (func(), error) { + if p == nil || p.isTripped() { + return nil, errRefreshSkipped + } + release := func() {} + if p.rateGate != nil { + var err error + release, err = p.rateGate.acquire(ctx) + if err != nil { + return nil, err + } + } + if p.isTripped() { + release() + return nil, errRefreshSkipped + } + return release, nil +} + +func (p *tokenRefreshProviderState) recordResult(err error) { + p.mu.Lock() + defer p.mu.Unlock() + + if err == nil { + p.consecutiveFailures = 0 + return + } + if errors.Is(err, errRefreshSkipped) || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + var attemptTimeoutErr *refreshAttemptTimeoutError + if !errors.As(err, &attemptTimeoutErr) { + return + } + } + var attemptTimeoutErr *refreshAttemptTimeoutError + if errors.As(err, &attemptTimeoutErr) { + p.consecutiveFailures++ + if p.consecutiveFailures >= p.service.providerFailureThreshold() { + p.tripped = true + } + return + } + var providerErr *providerConfigurationRefreshError + if errors.As(err, &providerErr) { + p.tripped = true + return + } + var containmentErr *providerCycleContainmentRefreshError + if errors.As(err, &containmentErr) { + p.tripped = true + return + } + var permanentErr *accountPermanentRefreshError + if errors.As(err, &permanentErr) { + p.consecutiveFailures = 0 + return + } + if isNonRetryableRefreshError(err) { + // A permanent account credential failure is isolated to that account and + // does not imply the provider is unhealthy. + p.consecutiveFailures = 0 + return + } + p.consecutiveFailures++ + if p.consecutiveFailures >= p.service.providerFailureThreshold() { + p.tripped = true + } +} + +// processRefresh preserves the existing test/internal call surface while the +// production loop supplies a cancelable parent context. +func (s *TokenRefreshService) processRefresh() { + s.processRefreshContext(context.Background()) +} + +// processRefreshContext executes one bounded, cursor-resumable refresh cycle. +func (s *TokenRefreshService) processRefreshContext(parent context.Context) { + if parent == nil { + parent = context.Background() + } + ctx, cancel := context.WithTimeout(parent, s.cycleTimeout()) + defer cancel() + + pager := s.candidatePager + if pager == nil { + pager, _ = s.accountRepo.(OAuthRefreshCandidatePager) + } + if pager == nil { + slog.Error("token_refresh.candidate_pager_missing") + return + } + platforms := s.eligiblePlatforms() + if len(platforms) == 0 { + slog.Error("token_refresh.provider_registry_empty") return } - totalAccounts := len(accounts) - oauthAccounts := 0 // 可刷新的OAuth账号数 - needsRefresh := 0 // 需要刷新的账号数 - refreshed, failed, skipped := 0, 0, 0 + refreshWindow := time.Duration(s.cfg.RefreshBeforeExpiryHours * float64(time.Hour)) + pageSize := s.candidatePageSize() + providerStates := make(map[string]*tokenRefreshProviderState, len(s.registrations)) + for i := range s.registrations { + registration := s.registrations[i] + providerStates[registration.platform] = &tokenRefreshProviderState{ + service: s, + registration: registration, + rateGate: s.providerRateGate(registration.platform), + poolGate: s.providerConcurrencyGate(registration.platform), + } + } - for i := range accounts { - account := &accounts[i] + stats := tokenRefreshPageStats{} + afterID := s.candidateAfterID() + for { + if ctx.Err() != nil { + slog.Warn("token_refresh.cycle_stopped", "error", ctx.Err(), "resume_after_id", afterID) + break + } + page, err := pager.ListOAuthRefreshCandidatePage(ctx, OAuthRefreshPageOptions{ + Platforms: platforms, + AfterID: afterID, + Limit: pageSize, + ActiveOnly: true, + IncludeSetupToken: true, + RequireRefreshToken: true, + ExcludeRetryCooldown: true, + }) + if err != nil { + slog.Error("token_refresh.list_accounts_failed", "error", err, "after_id", afterID) + break + } + if page == nil { + slog.Error("token_refresh.nil_candidate_page", "after_id", afterID) + break + } + accounts := page.Accounts + if !page.HasMore && page.NextAfterID == 0 && len(accounts) == 0 { + s.setCandidateAfterID(0) + break + } + if page.NextAfterID <= afterID { + slog.Error("token_refresh.invalid_candidate_page_metadata", "after_id", afterID) + break + } + if !isStrictlyIncreasingAccountPage(accounts, afterID) { + slog.Error("token_refresh.invalid_candidate_page", "after_id", afterID, "count", len(accounts)) + break + } - // 遍历所有刷新器，找到能处理此账号的 - for idx, refresher := range s.refreshers { - if !refresher.CanRefresh(account) { - continue - } + pageStats := s.processCandidatePage(ctx, accounts, providerStates, refreshWindow) + stats.total += pageStats.total + stats.oauth += pageStats.oauth + stats.needsRefresh += pageStats.needsRefresh + stats.refreshed += pageStats.refreshed + stats.skipped += pageStats.skipped + stats.failed += pageStats.failed - oauthAccounts++ - - // 检查是否需要刷新 - if !refresher.NeedsRefresh(account, refreshWindow) { - break // 不需要刷新，跳过 - } - - needsRefresh++ - - // 获取对应的 executor - var executor OAuthRefreshExecutor - if idx < len(s.executors) { - executor = s.executors[idx] - } - - // 执行刷新 - if err := s.refreshWithRetry(ctx, account, refresher, executor, refreshWindow); err != nil { - if errors.Is(err, errRefreshSkipped) { - skipped++ - } else { - slog.Warn("token_refresh.account_refresh_failed", - "account_id", account.ID, - "account_name", account.Name, - "error", err, - ) - failed++ - } - } else { - slog.Info("token_refresh.account_refreshed", - "account_id", account.ID, - "account_name", account.Name, - ) - refreshed++ - } - - // 每个账号只由一个refresher处理 + // Never advance past a partially processed page. Re-reading a page is + // safe because OAuthRefreshAPI re-reads DB state and checks expiry again. + if ctx.Err() != nil { + break + } + afterID = page.NextAfterID + s.setCandidateAfterID(afterID) + if !page.HasMore { + s.setCandidateAfterID(0) break } } - // 无刷新活动时降级为 Debug，有实际刷新活动时保持 Info - if needsRefresh == 0 && failed == 0 { + if stats.needsRefresh == 0 && stats.failed == 0 { slog.Debug("token_refresh.cycle_completed", - "total", totalAccounts, "oauth", oauthAccounts, - "needs_refresh", needsRefresh, "refreshed", refreshed, "skipped", skipped, "failed", failed) + "total", stats.total, "oauth", stats.oauth, + "needs_refresh", stats.needsRefresh, "refreshed", stats.refreshed, + "skipped", stats.skipped, "failed", stats.failed) } else { slog.Info("token_refresh.cycle_completed", - "total", totalAccounts, - "oauth", oauthAccounts, - "needs_refresh", needsRefresh, - "refreshed", refreshed, - "skipped", skipped, - "failed", failed, - ) + "total", stats.total, "oauth", stats.oauth, + "needs_refresh", stats.needsRefresh, "refreshed", stats.refreshed, + "skipped", stats.skipped, "failed", stats.failed) } } -// listActiveAccounts 获取后台 OAuth token 刷新候选账号。 -func (s *TokenRefreshService) listActiveAccounts(ctx context.Context) ([]Account, error) { - return s.accountRepo.ListOAuthRefreshCandidates(ctx) +func isStrictlyIncreasingAccountPage(accounts []Account, afterID int64) bool { + previous := afterID + for i := range accounts { + if accounts[i].ID <= previous { + return false + } + previous = accounts[i].ID + } + return true +} + +func (s *TokenRefreshService) processCandidatePage( + ctx context.Context, + accounts []Account, + providerStates map[string]*tokenRefreshProviderState, + refreshWindow time.Duration, +) tokenRefreshPageStats { + stats := tokenRefreshPageStats{total: len(accounts)} + groups := make(map[string][]*Account) + for i := range accounts { + account := &accounts[i] + state := providerStates[account.Platform] + if state == nil || state.registration.refresher == nil || !state.registration.refresher.CanRefresh(account) { + continue + } + stats.oauth++ + if !state.registration.refresher.NeedsRefresh(account, refreshWindow) { + continue + } + stats.needsRefresh++ + groups[account.Platform] = append(groups[account.Platform], account) + } + + type providerResult struct { + refreshed int + skipped int + failed int + } + results := make(chan providerResult, len(groups)) + var wg sync.WaitGroup + for platform, group := range groups { + state := providerStates[platform] + wg.Add(1) + go func() { + defer wg.Done() + refreshed, skipped, failed := s.processProviderAccounts(ctx, state, group, refreshWindow) + results <- providerResult{refreshed: refreshed, skipped: skipped, failed: failed} + }() + } + wg.Wait() + close(results) + for result := range results { + stats.refreshed += result.refreshed + stats.skipped += result.skipped + stats.failed += result.failed + } + return stats +} + +func (s *TokenRefreshService) processProviderAccounts( + ctx context.Context, + state *tokenRefreshProviderState, + accounts []*Account, + refreshWindow time.Duration, +) (refreshed, skipped, failed int) { + if state == nil || len(accounts) == 0 { + return 0, 0, 0 + } + type refreshResult struct { + accountID int64 + err error + } + jobs := make(chan *Account, len(accounts)) + results := make(chan refreshResult, len(accounts)) + workerCount := s.providerConcurrency() + if workerCount > len(accounts) { + workerCount = len(accounts) + } + var wg sync.WaitGroup + for i := 0; i < workerCount; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for account := range jobs { + if ctx.Err() != nil || state.isTripped() { + results <- refreshResult{accountID: account.ID, err: errRefreshSkipped} + continue + } + if state.isTripped() { + results <- refreshResult{accountID: account.ID, err: errRefreshSkipped} + continue + } + err := s.refreshWithRetryWithRateGate(ctx, account, state.registration.refresher, state.registration.executor, refreshWindow, state) + state.recordResult(err) + results <- refreshResult{accountID: account.ID, err: err} + } + }() + } + for _, account := range accounts { + jobs <- account + } + close(jobs) + wg.Wait() + close(results) + + for result := range results { + switch { + case result.err == nil: + refreshed++ + slog.Info("token_refresh.account_refreshed", "account_id", result.accountID, "platform", state.registration.platform) + case errors.Is(result.err, errRefreshSkipped): + skipped++ + default: + failed++ + slog.Warn("token_refresh.account_refresh_failed", "account_id", result.accountID, "platform", state.registration.platform, "error", logredact.RedactText(result.err.Error())) + } + } + return refreshed, skipped, failed +} + +func (s *TokenRefreshService) candidatePageSize() int { + if s.cfg != nil && s.cfg.CandidatePageSize > 0 { + return min(s.cfg.CandidatePageSize, maxTokenRefreshCandidatePageSize) + } + return defaultTokenRefreshCandidatePageSize +} + +func (s *TokenRefreshService) providerConcurrency() int { + if s.cfg != nil && s.cfg.ProviderConcurrency > 0 { + return min(s.cfg.ProviderConcurrency, maxTokenRefreshProviderConcurrency) + } + return defaultTokenRefreshProviderConcurrency +} + +func (s *TokenRefreshService) providerQPS() int { + if s.cfg != nil && s.cfg.ProviderQPS > 0 { + return min(s.cfg.ProviderQPS, maxTokenRefreshProviderQPS) + } + return defaultTokenRefreshProviderQPS +} + +// providerRateGate returns the process-local limiter shared by background +// cycles and admin reconciliation. Sharing it prevents concurrent entry points +// or retries from multiplying the configured per-provider request rate. +func (s *TokenRefreshService) providerRateGate(platform string) *tokenRefreshRateGate { + s.providerMu.Lock() + defer s.providerMu.Unlock() + if s.providerGates == nil { + s.providerGates = make(map[string]*tokenRefreshRateGate) + } + if gate := s.providerGates[platform]; gate != nil { + return gate + } + gate := newTokenRefreshRateGate(s.providerQPS()) + s.providerGates[platform] = gate + return gate +} + +// providerConcurrencyGate returns the process-local semaphore shared by every +// background cycle and admin reconciliation call for a provider. It is +// acquired and released around each upstream retry attempt, so parallel entry +// points cannot multiply ProviderConcurrency. +func (s *TokenRefreshService) providerConcurrencyGate(platform string) *tokenRefreshConcurrencyGate { + s.providerMu.Lock() + defer s.providerMu.Unlock() + if s.providerPools == nil { + s.providerPools = make(map[string]*tokenRefreshConcurrencyGate) + } + if gate := s.providerPools[platform]; gate != nil { + return gate + } + gate := newTokenRefreshConcurrencyGate(s.providerConcurrency()) + s.providerPools[platform] = gate + return gate +} + +func (s *TokenRefreshService) providerFailureThreshold() int { + if s.cfg != nil && s.cfg.ProviderFailureThreshold > 0 { + return min(s.cfg.ProviderFailureThreshold, maxTokenRefreshProviderFailureThreshold) + } + return defaultTokenRefreshProviderFailureThreshold +} + +func (s *TokenRefreshService) attemptTimeout() time.Duration { + timeout := defaultTokenRefreshAttemptTimeout + if s.attemptTimeoutOverride > 0 { + timeout = s.attemptTimeoutOverride + } else if s.cfg != nil && s.cfg.AttemptTimeoutSeconds > 0 { + seconds := min(s.cfg.AttemptTimeoutSeconds, int(maxTokenRefreshAttemptTimeout/time.Second)) + timeout = time.Duration(seconds) * time.Second + } + if s.refreshAPI != nil && s.refreshAPI.tokenCache != nil { + timeout = clampRefreshAttemptToLockLease(timeout, s.refreshAPI.lockTTL) + } + return timeout +} + +func clampRefreshAttemptToLockLease(timeout, lease time.Duration) time.Duration { + if timeout <= 0 || lease <= 0 { + return timeout + } + margin := lease / 10 + if margin > maxTokenRefreshLockSafetyMargin { + margin = maxTokenRefreshLockSafetyMargin + } + if margin <= 0 { + margin = time.Nanosecond + } + leaseBudget := lease - margin + if leaseBudget <= 0 { + leaseBudget = lease / 2 + } + if leaseBudget > 0 && timeout > leaseBudget { + return leaseBudget + } + return timeout +} + +func (s *TokenRefreshService) cycleTimeout() time.Duration { + if s.cfg != nil && s.cfg.CycleTimeoutSeconds > 0 { + seconds := min(s.cfg.CycleTimeoutSeconds, int(maxTokenRefreshCycleTimeout/time.Second)) + return time.Duration(seconds) * time.Second + } + return defaultTokenRefreshCycleTimeout +} + +func (s *TokenRefreshService) maxRetries() int { + if s.cfg != nil && s.cfg.MaxRetries > 0 { + return min(s.cfg.MaxRetries, maxTokenRefreshMaxRetries) + } + return defaultTokenRefreshMaxRetries } // refreshWithRetry 带重试的刷新 func (s *TokenRefreshService) refreshWithRetry(ctx context.Context, account *Account, refresher TokenRefresher, executor OAuthRefreshExecutor, refreshWindow time.Duration) error { - var lastErr error + return s.refreshWithRetryWithRateGate(ctx, account, refresher, executor, refreshWindow, nil) +} - for attempt := 1; attempt <= s.cfg.MaxRetries; attempt++ { +func (s *TokenRefreshService) refreshWithRetryWithRateGate( + ctx context.Context, + account *Account, + refresher TokenRefresher, + executor OAuthRefreshExecutor, + refreshWindow time.Duration, + gate refreshAttemptGate, +) error { + var lastErr error + maxRetries := s.maxRetries() + + for attempt := 1; attempt <= maxRetries; attempt++ { + if err := ctx.Err(); err != nil { + return err + } + releaseAttempt := func() {} + var acquireRate func(context.Context) (func(), error) + if gate != nil { + if providerGate, ok := gate.(providerRefreshAttemptGate); ok { + var err error + releaseAttempt, err = providerGate.acquire(ctx) + if err != nil { + return err + } + acquireRate = providerGate.acquireRate + } else { + // Compatibility gates are rate-admission gates. Acquire them only + // when an upstream Refresh call is actually about to start. + acquireRate = gate.acquire + } + } + attemptCtx, cancelAttempt := context.WithTimeout(ctx, s.attemptTimeout()) var newCredentials map[string]any var err error + shortCircuit := false + credentialsPersisted := false // 优先使用统一 API（带分布式锁 + DB 重读保护） if s.refreshAPI != nil && executor != nil { - result, refreshErr := s.refreshAPI.RefreshIfNeeded(ctx, account, executor, refreshWindow) + actualExecutor := executor + if acquireRate != nil { + actualExecutor = &rateLimitedOAuthRefreshExecutor{ + OAuthRefreshExecutor: executor, + acquireRate: acquireRate, + } + } + result, refreshErr := s.refreshAPI.RefreshIfNeeded(attemptCtx, account, actualExecutor, refreshWindow) + if result != nil && result.Account != nil { + account = result.Account + } if refreshErr != nil { err = refreshErr } else if result.LockHeld { // 锁被其他 worker 持有，由调用侧策略决定如何计数 - return s.refreshPolicy.handleLockHeld() + err = s.refreshPolicy.handleLockHeld() + shortCircuit = true } else if !result.Refreshed { // 已被其他路径刷新，由调用侧策略决定如何计数 - return s.refreshPolicy.handleAlreadyRefreshed() + err = s.refreshPolicy.handleAlreadyRefreshed() + shortCircuit = true } else { - account = result.Account + credentialsPersisted = result.NewCredentials != nil _ = result.NewCredentials // 统一 API 已设置 _token_version 并更新 DB，无需重复操作 } } else { // 降级：直接调用 refresher（兼容旧路径） - newCredentials, err = refresher.Refresh(ctx, account) - if newCredentials != nil { + releaseRate := func() {} + if acquireRate != nil { + releaseRate, err = acquireRate(attemptCtx) + } + if err == nil { + newCredentials, err = refresher.Refresh(attemptCtx, account) + } + if releaseRate != nil { + releaseRate() + } + attemptTimedOut := errors.Is(attemptCtx.Err(), context.DeadlineExceeded) && ctx.Err() == nil + if err == nil && newCredentials != nil && !attemptTimedOut { newCredentials["_token_version"] = time.Now().UnixMilli() - if saveErr := persistAccountCredentials(ctx, s.accountRepo, account, newCredentials); saveErr != nil { - return fmt.Errorf("failed to save credentials: %w", saveErr) + if saveErr := persistAccountCredentials(attemptCtx, s.accountRepo, account, newCredentials); saveErr != nil { + err = fmt.Errorf("failed to save credentials: %w", saveErr) + } else { + credentialsPersisted = true } } } + attemptTimedOut := errors.Is(attemptCtx.Err(), context.DeadlineExceeded) && ctx.Err() == nil + cancelAttempt() + releaseAttempt() + persistedAfterAttemptDeadline := attemptTimedOut && credentialsPersisted && err == nil + if attemptTimedOut && !persistedAfterAttemptDeadline && !isProviderScopedTerminalRefreshError(err) { + cause := err + if cause == nil { + cause = context.DeadlineExceeded + } + err = &refreshAttemptTimeoutError{err: cause} + shortCircuit = false + if credentialsPersisted { + s.postRefreshStateSyncWithCleanup(ctx, account) + } + } + if shortCircuit { + return err + } if err == nil { + if ctxErr := ctx.Err(); ctxErr != nil { + if credentialsPersisted { + s.postRefreshStateSyncWithCleanup(ctx, account) + } + return ctxErr + } + if persistedAfterAttemptDeadline { + // The provider result and exact-state CAS are already durable. Only + // the internal attempt budget elapsed while bounded detached cleanup + // completed; do not convert that success into retry/cooldown/breaker + // evidence. Publish cache state with a fresh cleanup context and stop. + s.postRefreshStateSyncWithCleanup(ctx, account) + return nil + } s.postRefreshActions(ctx, account) return nil } + if ctxErr := ctx.Err(); ctxErr != nil { + if credentialsPersisted { + s.postRefreshStateSyncWithCleanup(ctx, account) + } + return ctxErr + } + if errors.Is(err, errRefreshSkipped) { + return errRefreshSkipped + } + if isProviderScopedTerminalRefreshError(err) { + return err + } + var stateUnavailableErr *oauthRefreshStateUnavailableError + if errors.As(err, &stateUnavailableErr) { + return &providerCycleContainmentRefreshError{err: err} + } + if isAmbiguousGrokEntitlementRefreshError(account, err) { + // The current Grok client labels every token-endpoint 403 as an + // entitlement denial. Without explicit entitlement evidence, contain + // the provider for this cycle instead of disabling an account on a + // possible WAF or shared provider failure. + return &providerCycleContainmentRefreshError{err: err} + } + + // Provider-wide OAuth client/scope failures are not evidence that every + // account is invalid. Return a typed internal signal so the cycle contains + // the provider without mutating account state. + if isSharedProviderRefreshError(err) { + return &providerConfigurationRefreshError{err: err} + } // 不可重试错误（invalid_grant/invalid_client 等）直接标记 error 状态并返回 if isNonRetryableRefreshError(err) { errorMsg := "Token refresh failed (non-retryable): " + logredact.RedactText(err.Error()) - s.notifyAccountSchedulingBlocked(account, time.Time{}, "token_refresh_non_retryable") + isGrokOAuth := account.IsGrokOAuth() + if !isGrokOAuth { + s.notifyAccountSchedulingBlocked(account, time.Time{}, "token_refresh_non_retryable") + } s.clearAntigravityForceTokenRefresh(ctx, account, "non_retryable") - if setErr := s.accountRepo.SetError(ctx, account.ID, errorMsg); setErr != nil { + persistentlyBlocked := false + var setErr error + if isGrokOAuth { + conditionalRepo, ok := s.accountRepo.(GrokOAuthRefreshMutationRepository) + if !ok { + return &providerConfigurationRefreshError{ + err: errors.New("grok OAuth conditional refresh mutation repository is not configured"), + } + } else { + persistentlyBlocked, setErr = conditionalRepo.SetGrokOAuthRefreshErrorIfCredentialsUnchanged( + ctx, + account.ID, + account.Credentials, + account.ProxyID, + errorMsg, + ) + if setErr == nil && !persistentlyBlocked { + slog.Info("token_refresh.grok_error_status_skipped_stale_credentials", "account_id", account.ID) + return errRefreshSkipped + } + } + } else { + setErr = s.accountRepo.SetError(ctx, account.ID, errorMsg) + persistentlyBlocked = setErr == nil + } + if setErr != nil { slog.Error("token_refresh.set_error_status_failed", "account_id", account.ID, "error", setErr, ) + if isGrokOAuth { + return &providerCycleContainmentRefreshError{ + err: fmt.Errorf("failed to conditionally persist Grok OAuth refresh failure: %w", setErr), + } + } + } else if isGrokOAuth && persistentlyBlocked { + s.notifyAccountSchedulingBlocked(account, time.Time{}, "token_refresh_non_retryable") + } + cacheInvalidationFailed := false + if account.Type == AccountTypeOAuth && (!isGrokOAuth || persistentlyBlocked) { + if s.cacheInvalidator == nil { + cacheInvalidationFailed = true + } else if invalidateErr := s.cacheInvalidator.InvalidateToken(ctx, account); invalidateErr != nil { + cacheInvalidationFailed = true + slog.Warn("token_refresh.invalidate_failed_token_cache_failed", + "account_id", account.ID, + "error", logredact.RedactText(invalidateErr.Error()), + ) + } + } + return &accountPermanentRefreshError{ + err: err, + persistentlyBlocked: persistentlyBlocked, + cacheInvalidationFailed: cacheInvalidationFailed, } - return err } lastErr = err slog.Warn("token_refresh.retry_attempt_failed", "account_id", account.ID, "attempt", attempt, - "max_retries", s.cfg.MaxRetries, - "error", err, + "max_retries", maxRetries, + "error", logredact.RedactText(err.Error()), ) // 如果还有重试机会，等待后重试 - if attempt < s.cfg.MaxRetries { - // 指数退避：2^(attempt-1) * baseSeconds - backoff := time.Duration(s.cfg.RetryBackoffSeconds) * time.Second * time.Duration(1<<(attempt-1)) - time.Sleep(backoff) + if attempt < maxRetries { + backoff := s.retryBackoff(account.ID, attempt) + if backoff > 0 { + timer := time.NewTimer(backoff) + select { + case <-ctx.Done(): + timer.Stop() + return ctx.Err() + case <-timer.C: + } + } } } + if err := ctx.Err(); err != nil { + return err + } // 可重试错误耗尽：临时标记账号不可调度，避免请求路径反复命中已知失败的账号 slog.Warn("token_refresh.retry_exhausted", "account_id", account.ID, "platform", account.Platform, - "max_retries", s.cfg.MaxRetries, - "error", lastErr, + "max_retries", maxRetries, + "error", logredact.RedactText(lastErr.Error()), ) // 设置临时不可调度 10 分钟（不标记 error，保持 status=active 让下个刷新周期能继续尝试） @@ -352,6 +1087,42 @@ func (s *TokenRefreshService) refreshWithRetry(ctx context.Context, account *Acc if lastErr != nil { reason += ": " + logredact.RedactText(lastErr.Error()) } + if account.IsGrokOAuth() { + conditionalRepo, ok := s.accountRepo.(GrokOAuthRefreshMutationRepository) + if !ok { + return &providerConfigurationRefreshError{ + err: errors.New("grok OAuth conditional refresh mutation repository is not configured"), + } + } + applied, setErr := conditionalRepo.SetGrokOAuthRefreshTempUnschedulableIfCredentialsUnchanged( + ctx, + account.ID, + account.Credentials, + account.ProxyID, + until, + reason, + ) + if setErr != nil { + slog.Warn("token_refresh.set_temp_unschedulable_failed", + "account_id", account.ID, + "error", setErr, + ) + return &providerCycleContainmentRefreshError{ + err: fmt.Errorf("failed to conditionally persist Grok OAuth refresh cooldown: %w", setErr), + } + } else if !applied { + slog.Info("token_refresh.grok_temp_unschedulable_skipped_stale_credentials", "account_id", account.ID) + return errRefreshSkipped + } else { + s.notifyAccountSchedulingBlocked(account, until, "token_refresh_retry_exhausted") + slog.Info("token_refresh.temp_unschedulable_set", + "account_id", account.ID, + "until", until.Format(time.RFC3339), + ) + } + return lastErr + } + s.notifyAccountSchedulingBlocked(account, until, "token_refresh_retry_exhausted") if setErr := s.accountRepo.SetTempUnschedulable(ctx, account.ID, until, reason); setErr != nil { slog.Warn("token_refresh.set_temp_unschedulable_failed", @@ -368,6 +1139,23 @@ func (s *TokenRefreshService) refreshWithRetry(ctx context.Context, account *Acc return lastErr } +func (s *TokenRefreshService) retryBackoff(accountID int64, attempt int) time.Duration { + if s.cfg == nil || s.cfg.RetryBackoffSeconds <= 0 { + return 0 + } + shift := attempt - 1 + if shift > 10 { + shift = 10 + } + baseSeconds := min(s.cfg.RetryBackoffSeconds, int(maxTokenRefreshRetryBackoff/time.Second)) + base := time.Duration(baseSeconds) * time.Second * time.Duration(1<= 0 { + body := msg[bodyIndex+len("body:"):] + for _, evidence := range []string{ + "entitlement_denied", + "entitlement denied", + "subscription_required", + "no_active_subscription", + } { + if strings.Contains(body, evidence) { + return false + } + } + } + return true +} + +func (e *providerConfigurationRefreshError) Unwrap() error { + if e == nil { + return nil + } + return e.err +} + +func isSharedProviderRefreshError(err error) bool { + if err == nil { + return false + } + msg := strings.ToLower(err.Error()) + for _, needle := range []string{ + "invalid_client", + "unauthorized_client", + "invalid_scope", + "unknown scope", + } { + if strings.Contains(msg, needle) { + return true + } + } + return false +} + // isNonRetryableRefreshError 判断是否为不可重试的刷新错误 // 这些错误通常表示凭证已失效或配置确实缺失，需要用户重新授权 // 注意：missing_project_id 错误只在真正缺失（从未获取过）时返回，临时获取失败不会返回此错误 diff --git a/backend/internal/service/token_refresh_service_candidates_test.go b/backend/internal/service/token_refresh_service_candidates_test.go index 5e2c8cc4aa..ac4e3cf4f0 100644 --- a/backend/internal/service/token_refresh_service_candidates_test.go +++ b/backend/internal/service/token_refresh_service_candidates_test.go @@ -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") } diff --git a/backend/internal/service/token_refresh_service_test.go b/backend/internal/service/token_refresh_service_test.go index c4c0b35ceb..75a84bfbcc 100644 --- a/backend/internal/service/token_refresh_service_test.go +++ b/backend/internal/service/token_refresh_service_test.go @@ -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 的 service（Path 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{ diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go index a8674382f2..e4318e6079 100644 --- a/backend/internal/service/wire.go +++ b/backend/internal/service/wire.go @@ -687,6 +687,7 @@ var ProviderSet = wire.NewSet( NewCRSSyncService, ProvideUpdateService, ProvideTokenRefreshService, + wire.Bind(new(GrokOAuthReconciler), new(*TokenRefreshService)), ProvideAccountExpiryService, ProvideProxyExpiryService, ProvideSubscriptionExpiryService, diff --git a/deploy/config.example.yaml b/deploy/config.example.yaml index 7b98993c53..93fa6a2b59 100644 --- a/deploy/config.example.yaml +++ b/deploy/config.example.yaml @@ -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