diff --git a/backend/internal/repository/account_repo.go b/backend/internal/repository/account_repo.go index e7626696bf..fb59de6717 100644 --- a/backend/internal/repository/account_repo.go +++ b/backend/internal/repository/account_repo.go @@ -649,6 +649,58 @@ func (r *accountRepository) ListActive(ctx context.Context) ([]service.Account, return r.accountsToService(ctx, accounts) } +func (r *accountRepository) ListOAuthRefreshCandidates(ctx context.Context) ([]service.Account, error) { + if r.sql == nil { + return nil, errors.New("account repository SQL executor not configured") + } + rows, err := r.sql.QueryContext(ctx, ` + SELECT id + FROM accounts + WHERE deleted_at IS NULL + AND status = 'active' + AND type = 'oauth' + AND platform IN ('anthropic', 'openai', 'gemini', 'antigravity') + AND credentials ? 'refresh_token' + AND btrim(credentials->>'refresh_token') <> '' + AND NOT ( + temp_unschedulable_until > NOW() + AND temp_unschedulable_reason LIKE 'token refresh retry exhausted:%' + ) + ORDER BY priority ASC, id ASC + `) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + + var ids []int64 + for rows.Next() { + var id int64 + if err := rows.Scan(&id); err != nil { + return nil, err + } + ids = append(ids, id) + } + if err := rows.Err(); err != nil { + return nil, err + } + if len(ids) == 0 { + return []service.Account{}, nil + } + + accounts, err := r.GetByIDs(ctx, ids) + if err != nil { + return nil, err + } + out := make([]service.Account, 0, len(accounts)) + for _, account := range accounts { + if account != nil { + out = append(out, *account) + } + } + return out, nil +} + func (r *accountRepository) ListByPlatform(ctx context.Context, platform string) ([]service.Account, error) { accounts, err := r.client.Account.Query(). Where( @@ -1148,7 +1200,7 @@ func (r *accountRepository) SetOverloaded(ctx context.Context, id int64, until t } func (r *accountRepository) SetTempUnschedulable(ctx context.Context, id int64, until time.Time, reason string) error { - _, err := r.sql.ExecContext(ctx, ` + result, err := r.sql.ExecContext(ctx, ` UPDATE accounts SET temp_unschedulable_until = $1, temp_unschedulable_reason = $2, @@ -1160,6 +1212,13 @@ func (r *accountRepository) SetTempUnschedulable(ctx context.Context, id int64, if err != nil { return err } + affected, err := result.RowsAffected() + if err != nil { + return err + } + if affected <= 0 { + return nil + } if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil { logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue temp unschedulable failed: account=%d err=%v", id, err) } diff --git a/backend/internal/repository/account_repo_temp_unsched_test.go b/backend/internal/repository/account_repo_temp_unsched_test.go new file mode 100644 index 0000000000..e802e3b0cb --- /dev/null +++ b/backend/internal/repository/account_repo_temp_unsched_test.go @@ -0,0 +1,98 @@ +package repository + +import ( + "context" + "database/sql" + "regexp" + "strings" + "testing" + "time" + + sqlmock "github.com/DATA-DOG/go-sqlmock" + "github.com/stretchr/testify/require" +) + +func TestAccountRepository_SetTempUnschedulable_NoRowsAffectedDoesNotWriteOutbox(t *testing.T) { + exec := &recordingSQLExecutor{result: rowsAffectedResult(0)} + repo := newAccountRepositoryWithSQL(nil, exec, nil) + until := time.Now().Add(10 * time.Minute) + + err := repo.SetTempUnschedulable(context.Background(), 42, until, "retry") + require.NoError(t, err) + require.Len(t, exec.execQueries, 1) + require.Contains(t, exec.execQueries[0], "UPDATE accounts") + require.NotContains(t, strings.Join(exec.execQueries, "\n"), "scheduler_outbox") +} + +func TestAccountRepository_ListOAuthRefreshCandidates_SQLFilter(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"})). + WillDelayFor(0) + + repo := newAccountRepositoryWithSQL(nil, captureQuerySQL{db: db, captured: &capturedSQL}, nil) + + accounts, err := repo.ListOAuthRefreshCandidates(context.Background()) + require.NoError(t, err) + require.Empty(t, accounts) + + normalized := normalizeSQLWhitespace(capturedSQL) + require.Contains(t, normalized, "deleted_at IS NULL") + require.Contains(t, normalized, "status = 'active'") + require.Contains(t, normalized, "type = 'oauth'") + require.Contains(t, normalized, "platform IN ('anthropic', 'openai', 'gemini', 'antigravity')") + require.Contains(t, normalized, "credentials ? 'refresh_token'") + require.Contains(t, normalized, "btrim(credentials->>'refresh_token') <> ''") + require.Contains(t, normalized, "temp_unschedulable_until > NOW()") + require.Contains(t, normalized, "temp_unschedulable_reason LIKE 'token refresh retry exhausted:%'") + require.Contains(t, normalized, "ORDER BY priority ASC, id ASC") + require.NotContains(t, normalized, "credentials->>'expires_at'") + require.NoError(t, mock.ExpectationsWereMet()) +} + +type captureQuerySQL struct { + db *sql.DB + captured *string +} + +func (c captureQuerySQL) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) { + return c.db.ExecContext(ctx, query, args...) +} + +func (c captureQuerySQL) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) { + if c.captured != nil { + *c.captured = query + } + return c.db.QueryContext(ctx, query, args...) +} + +func normalizeSQLWhitespace(sql string) string { + return strings.Join(regexp.MustCompile(`\s+`).Split(strings.TrimSpace(sql), -1), " ") +} + +type rowsAffectedResult int64 + +func (r rowsAffectedResult) LastInsertId() (int64, error) { return 0, nil } +func (r rowsAffectedResult) RowsAffected() (int64, error) { return int64(r), nil } + +type recordingSQLExecutor struct { + result sql.Result + err error + execQueries []string +} + +func (e *recordingSQLExecutor) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) { + e.execQueries = append(e.execQueries, query) + if e.err != nil { + return nil, e.err + } + return e.result, nil +} + +func (e *recordingSQLExecutor) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) { + return nil, sql.ErrNoRows +} diff --git a/backend/internal/service/account_repository_refresh_candidates_unit_test.go b/backend/internal/service/account_repository_refresh_candidates_unit_test.go new file mode 100644 index 0000000000..7dbf571f7d --- /dev/null +++ b/backend/internal/service/account_repository_refresh_candidates_unit_test.go @@ -0,0 +1,25 @@ +//go:build unit + +package service + +import "context" + +func (s *accountRepoStub) ListOAuthRefreshCandidates(context.Context) ([]Account, error) { + panic("unexpected ListOAuthRefreshCandidates call") +} + +func (r *openAIAccountTestRepo) ListOAuthRefreshCandidates(context.Context) ([]Account, error) { + panic("unexpected ListOAuthRefreshCandidates call") +} + +func (m *groupAwareMockAccountRepo) ListOAuthRefreshCandidates(context.Context) ([]Account, error) { + panic("unexpected ListOAuthRefreshCandidates call") +} + +func (m *mockAccountRepoForPlatform) ListOAuthRefreshCandidates(context.Context) ([]Account, error) { + panic("unexpected ListOAuthRefreshCandidates call") +} + +func (m *mockAccountRepoForGemini) ListOAuthRefreshCandidates(context.Context) ([]Account, error) { + return m.ListActive(context.Background()) +} diff --git a/backend/internal/service/account_service.go b/backend/internal/service/account_service.go index 27a0ff8946..41b5a2d97e 100644 --- a/backend/internal/service/account_service.go +++ b/backend/internal/service/account_service.go @@ -41,6 +41,7 @@ type AccountRepository interface { ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, *pagination.PaginationResult, 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/ratelimit_session_window_test.go b/backend/internal/service/ratelimit_session_window_test.go index 0c2d711228..d2ff0d835c 100644 --- a/backend/internal/service/ratelimit_session_window_test.go +++ b/backend/internal/service/ratelimit_session_window_test.go @@ -93,6 +93,9 @@ func (m *sessionWindowMockRepo) ListByGroup(context.Context, int64) ([]Account, func (m *sessionWindowMockRepo) ListActive(context.Context) ([]Account, error) { panic("unexpected") } +func (m *sessionWindowMockRepo) ListOAuthRefreshCandidates(context.Context) ([]Account, error) { + panic("unexpected") +} func (m *sessionWindowMockRepo) ListByPlatform(context.Context, string) ([]Account, error) { panic("unexpected") } diff --git a/backend/internal/service/token_refresh_service.go b/backend/internal/service/token_refresh_service.go index bac7e4cbed..7c66faea89 100644 --- a/backend/internal/service/token_refresh_service.go +++ b/backend/internal/service/token_refresh_service.go @@ -12,8 +12,12 @@ import ( "github.com/Wei-Shaw/sub2api/internal/config" ) -// tokenRefreshTempUnschedDuration token 刷新重试耗尽后临时不可调度的持续时间 -const tokenRefreshTempUnschedDuration = 10 * time.Minute +const ( + // tokenRefreshTempUnschedDuration token 刷新重试耗尽后临时不可调度的持续时间 + tokenRefreshTempUnschedDuration = 10 * time.Minute + + tokenRefreshRetryExhaustedReasonPrefix = "token refresh retry exhausted:" +) // TokenRefreshService OAuth token自动刷新服务 // 定期检查并刷新即将过期的token @@ -255,10 +259,9 @@ func (s *TokenRefreshService) processRefresh() { } } -// listActiveAccounts 获取所有active状态的账号 -// 使用ListActive确保刷新所有活跃账号的token（包括临时禁用的） +// listActiveAccounts 获取后台 OAuth token 刷新候选账号。 func (s *TokenRefreshService) listActiveAccounts(ctx context.Context) ([]Account, error) { - return s.accountRepo.ListActive(ctx) + return s.accountRepo.ListOAuthRefreshCandidates(ctx) } // refreshWithRetry 带重试的刷新 @@ -310,9 +313,6 @@ func (s *TokenRefreshService) refreshWithRetry(ctx context.Context, account *Acc "error", setErr, ) } - // 刷新失败但 access_token 可能仍有效，尝试设置隐私 - s.ensureOpenAIPrivacy(ctx, account) - s.ensureAntigravityPrivacy(ctx, account) return err } @@ -340,13 +340,9 @@ func (s *TokenRefreshService) refreshWithRetry(ctx context.Context, account *Acc "error", lastErr, ) - // 刷新失败但 access_token 可能仍有效，尝试设置隐私 - s.ensureOpenAIPrivacy(ctx, account) - s.ensureAntigravityPrivacy(ctx, account) - // 设置临时不可调度 10 分钟（不标记 error，保持 status=active 让下个刷新周期能继续尝试） until := time.Now().Add(tokenRefreshTempUnschedDuration) - reason := fmt.Sprintf("token refresh retry exhausted: %v", lastErr) + reason := fmt.Sprintf("%s %v", tokenRefreshRetryExhaustedReasonPrefix, 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", diff --git a/backend/internal/service/token_refresh_service_candidates_test.go b/backend/internal/service/token_refresh_service_candidates_test.go new file mode 100644 index 0000000000..fc5fc9ddbc --- /dev/null +++ b/backend/internal/service/token_refresh_service_candidates_test.go @@ -0,0 +1,190 @@ +package service + +import ( + "context" + "errors" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/imroc/req/v3" + "github.com/stretchr/testify/require" +) + +type tokenRefreshCandidateRepo struct { + AccountRepository + accounts []Account + updatedCredentialIDs []int64 + setErrorCalls int + setTempUnschedCalls int + lastTempUnschedReason string + listActiveCalls int +} + +func (r *tokenRefreshCandidateRepo) ListActive(context.Context) ([]Account, error) { + r.listActiveCalls++ + return r.accounts, nil +} + +func (r *tokenRefreshCandidateRepo) ListOAuthRefreshCandidates(context.Context) ([]Account, error) { + candidates := make([]Account, 0, len(r.accounts)) + now := time.Now() + for _, account := range r.accounts { + refreshToken, _ := account.Credentials["refresh_token"].(string) + inRetryCooldown := account.TempUnschedulableUntil != nil && + account.TempUnschedulableUntil.After(now) && + strings.HasPrefix(account.TempUnschedulableReason, tokenRefreshRetryExhaustedReasonPrefix) + if account.Status != StatusActive || + account.Type != AccountTypeOAuth || + !isOAuthRefreshPlatform(account.Platform) || + strings.TrimSpace(refreshToken) == "" || + inRetryCooldown { + continue + } + candidates = append(candidates, account) + } + return candidates, nil +} + +func (r *tokenRefreshCandidateRepo) UpdateCredentials(_ context.Context, id int64, _ map[string]any) error { + r.updatedCredentialIDs = append(r.updatedCredentialIDs, id) + return nil +} + +func (r *tokenRefreshCandidateRepo) SetError(context.Context, int64, string) error { + r.setErrorCalls++ + return nil +} + +func (r *tokenRefreshCandidateRepo) SetTempUnschedulable(_ context.Context, _ int64, _ time.Time, reason string) error { + r.setTempUnschedCalls++ + r.lastTempUnschedReason = reason + return nil +} + +func isOAuthRefreshPlatform(platform string) bool { + switch platform { + case PlatformAnthropic, PlatformOpenAI, PlatformGemini, PlatformAntigravity: + return true + default: + return false + } +} + +type tokenRefreshTestRefresher struct { + err error +} + +func (r *tokenRefreshTestRefresher) CanRefresh(*Account) bool { return true } + +func (r *tokenRefreshTestRefresher) NeedsRefresh(*Account, time.Duration) bool { return true } + +func (r *tokenRefreshTestRefresher) Refresh(context.Context, *Account) (map[string]any, error) { + if r.err != nil { + return nil, r.err + } + return map[string]any{"access_token": "new-access-token", "refresh_token": "new-refresh-token"}, nil +} + +func TestTokenRefreshService_ProcessRefreshUsesOAuthRefreshCandidates(t *testing.T) { + future := time.Now().Add(10 * time.Minute) + repo := &tokenRefreshCandidateRepo{ + accounts: []Account{ + { + ID: 1, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Credentials: map[string]any{"refresh_token": "refresh-token"}, + }, + { + ID: 2, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Credentials: map[string]any{}, + }, + { + ID: 3, + Platform: PlatformGemini, + Type: AccountTypeAPIKey, + Status: StatusActive, + Credentials: map[string]any{"refresh_token": "refresh-token"}, + }, + { + ID: 4, + Platform: PlatformAntigravity, + Type: AccountTypeOAuth, + Status: StatusActive, + Credentials: map[string]any{"refresh_token": "refresh-token"}, + TempUnschedulableUntil: &future, + TempUnschedulableReason: tokenRefreshRetryExhaustedReasonPrefix + " network timeout", + }, + { + ID: 5, + Platform: "other", + Type: AccountTypeOAuth, + Status: StatusActive, + Credentials: map[string]any{"refresh_token": "refresh-token"}, + }, + }, + } + svc := &TokenRefreshService{ + accountRepo: repo, + refreshers: []TokenRefresher{&tokenRefreshTestRefresher{}}, + refreshPolicy: DefaultBackgroundRefreshPolicy(), + cfg: &config.TokenRefreshConfig{RefreshBeforeExpiryHours: 1, MaxRetries: 1}, + } + + svc.processRefresh() + + require.Zero(t, repo.listActiveCalls, "TokenRefreshService should not use the broad active-account query") + require.Equal(t, []int64{1}, repo.updatedCredentialIDs) +} + +func TestTokenRefreshService_RefreshFailureDoesNotCallPrivacy(t *testing.T) { + tests := []struct { + name string + err error + }{ + {name: "retry exhausted", err: errors.New("temporary upstream timeout")}, + {name: "non retryable", err: errors.New("invalid_grant: token revoked")}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + repo := &tokenRefreshCandidateRepo{} + svc := &TokenRefreshService{ + accountRepo: repo, + refreshPolicy: DefaultBackgroundRefreshPolicy(), + cfg: &config.TokenRefreshConfig{MaxRetries: 1, RetryBackoffSeconds: 0}, + privacyClientFactory: func(string) (*req.Client, error) { + t.Fatalf("privacy client factory must not be called on refresh failure") + return nil, errors.New("unexpected privacy call") + }, + } + account := &Account{ + ID: 11, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "access_token": "old-access-token", + "refresh_token": "refresh-token", + }, + } + + err := svc.refreshWithRetry(context.Background(), account, &tokenRefreshTestRefresher{err: tt.err}, nil, time.Hour) + + require.Error(t, err) + if isNonRetryableRefreshError(tt.err) { + require.Equal(t, 1, repo.setErrorCalls) + require.Zero(t, repo.setTempUnschedCalls) + } else { + require.Zero(t, repo.setErrorCalls) + require.Equal(t, 1, repo.setTempUnschedCalls) + require.True(t, strings.HasPrefix(repo.lastTempUnschedReason, tokenRefreshRetryExhaustedReasonPrefix)) + } + }) + } +}