wanwu/sub2api
1
0
Fork 0
mirror of https://gitee.com/wanwujie/sub2api synced 2026-04-11 02:24:45 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #854 from james-6-23/main

Browse Source 
feat(admin): 支持定时测试自动恢复并统一账号恢复入口
This commit is contained in:
Wesley Liddick
2026-03-09 08:48:36 +08:00
committed by GitHub
parent 00c151b463 34aad82ac3
commit faf6441633
23 changed files with 535 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
backend/internal/service/admin_service.go
View File

@@ -1723,16 +1723,10 @@ func (s *adminServiceImpl) RefreshAccountCredentials(ctx context.Context, id int
}
func (s *adminServiceImpl) ClearAccountError(ctx context.Context, id int64) (*Account, error) {
account, err := s.accountRepo.GetByID(ctx, id)
if err != nil {
if err := s.accountRepo.ClearError(ctx, id); err != nil {
return nil, err
}
account.Status = StatusActive
account.ErrorMessage = ""
if err := s.accountRepo.Update(ctx, account); err != nil {
return nil, err
}
return account, nil
return s.accountRepo.GetByID(ctx, id)
}
func (s *adminServiceImpl) SetAccountError(ctx context.Context, id int64, errorMsg string) error {

77
backend/internal/service/ratelimit_service.go
View File

@@ -28,6 +28,17 @@ type RateLimitService struct {
usageCache map[int64]*geminiUsageCacheEntry
}
// SuccessfulTestRecoveryResult 表示测试成功后恢复了哪些运行时状态。
type SuccessfulTestRecoveryResult struct {
ClearedError bool
ClearedRateLimit bool
}
// AccountRecoveryOptions 控制账号恢复时的附加行为。
type AccountRecoveryOptions struct {
InvalidateToken bool
}
type geminiUsageCacheEntry struct {
windowStart time.Time
cachedAt time.Time
@@ -1040,6 +1051,42 @@ func (s *RateLimitService) ClearRateLimit(ctx context.Context, accountID int64)
return nil
}
// RecoverAccountState 按需恢复账号的可恢复运行时状态。
func (s *RateLimitService) RecoverAccountState(ctx context.Context, accountID int64, options AccountRecoveryOptions) (*SuccessfulTestRecoveryResult, error) {
account, err := s.accountRepo.GetByID(ctx, accountID)
if err != nil {
return nil, err
}
result := &SuccessfulTestRecoveryResult{}
if account.Status == StatusError {
if err := s.accountRepo.ClearError(ctx, accountID); err != nil {
return nil, err
}
result.ClearedError = true
if options.InvalidateToken && s.tokenCacheInvalidator != nil && account.IsOAuth() {
if invalidateErr := s.tokenCacheInvalidator.InvalidateToken(ctx, account); invalidateErr != nil {
slog.Warn("recover_account_state_invalidate_token_failed", "account_id", accountID, "error", invalidateErr)
}
}
}
if hasRecoverableRuntimeState(account) {
if err := s.ClearRateLimit(ctx, accountID); err != nil {
return nil, err
}
result.ClearedRateLimit = true
}
return result, nil
}
// RecoverAccountAfterSuccessfulTest 将一次成功测试视为正常请求，
// 按需恢复 error / rate-limit / overload / temp-unsched / model-rate-limit 等运行时状态。
func (s *RateLimitService) RecoverAccountAfterSuccessfulTest(ctx context.Context, accountID int64) (*SuccessfulTestRecoveryResult, error) {
return s.RecoverAccountState(ctx, accountID, AccountRecoveryOptions{})
}
func (s *RateLimitService) ClearTempUnschedulable(ctx context.Context, accountID int64) error {
if err := s.accountRepo.ClearTempUnschedulable(ctx, accountID); err != nil {
return err
@@ -1056,6 +1103,36 @@ func (s *RateLimitService) ClearTempUnschedulable(ctx context.Context, accountID
return nil
}
func hasRecoverableRuntimeState(account *Account) bool {
if account == nil {
return false
}
if account.RateLimitedAt != nil || account.RateLimitResetAt != nil || account.OverloadUntil != nil || account.TempUnschedulableUntil != nil {
return true
}
if len(account.Extra) == 0 {
return false
}
return hasNonEmptyMapValue(account.Extra, "model_rate_limits") || hasNonEmptyMapValue(account.Extra, "antigravity_quota_scopes")
}
func hasNonEmptyMapValue(extra map[string]any, key string) bool {
raw, ok := extra[key]
if !ok || raw == nil {
return false
}
switch typed := raw.(type) {
case map[string]any:
return len(typed) > 0
case map[string]string:
return len(typed) > 0
case []any:
return len(typed) > 0
default:
return true
}
}
func (s *RateLimitService) GetTempUnschedStatus(ctx context.Context, accountID int64) (*TempUnschedState, error) {
now := time.Now().Unix()
if s.tempUnschedCache != nil {

134
backend/internal/service/ratelimit_service_clear_test.go
View File

@@ -6,6 +6,7 @@ import (
"context"
"errors"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/stretchr/testify/require"
@@ -13,16 +14,34 @@ import (
type rateLimitClearRepoStub struct {
mockAccountRepoForGemini
getByIDAccount *Account
getByIDErr error
getByIDCalls int
clearErrorCalls int
clearRateLimitCalls int
clearAntigravityCalls int
clearModelRateLimitCalls int
clearTempUnschedCalls int
clearErrorErr error
clearRateLimitErr error
clearAntigravityErr error
clearModelRateLimitErr error
clearTempUnschedulableErr error
}
func (r *rateLimitClearRepoStub) GetByID(ctx context.Context, id int64) (*Account, error) {
r.getByIDCalls++
if r.getByIDErr != nil {
return nil, r.getByIDErr
}
return r.getByIDAccount, nil
}
func (r *rateLimitClearRepoStub) ClearError(ctx context.Context, id int64) error {
r.clearErrorCalls++
return r.clearErrorErr
}
func (r *rateLimitClearRepoStub) ClearRateLimit(ctx context.Context, id int64) error {
r.clearRateLimitCalls++
return r.clearRateLimitErr
@@ -48,6 +67,11 @@ type tempUnschedCacheRecorder struct {
deleteErr error
}
type recoverTokenInvalidatorStub struct {
accounts []*Account
err error
}
func (c *tempUnschedCacheRecorder) SetTempUnsched(ctx context.Context, accountID int64, state *TempUnschedState) error {
return nil
}
@@ -61,6 +85,11 @@ func (c *tempUnschedCacheRecorder) DeleteTempUnsched(ctx context.Context, accoun
return c.deleteErr
}
func (s *recoverTokenInvalidatorStub) InvalidateToken(ctx context.Context, account *Account) error {
s.accounts = append(s.accounts, account)
return s.err
}
func TestRateLimitService_ClearRateLimit_AlsoClearsTempUnschedulable(t *testing.T) {
repo := &rateLimitClearRepoStub{}
cache := &tempUnschedCacheRecorder{}
@@ -170,3 +199,108 @@ func TestRateLimitService_ClearRateLimit_WithoutTempUnschedCache(t *testing.T) {
require.Equal(t, 1, repo.clearModelRateLimitCalls)
require.Equal(t, 1, repo.clearTempUnschedCalls)
}
func TestRateLimitService_RecoverAccountAfterSuccessfulTest_ClearsErrorAndRateLimitRelatedState(t *testing.T) {
now := time.Now()
repo := &rateLimitClearRepoStub{
getByIDAccount: &Account{
ID: 42,
Status: StatusError,
RateLimitedAt: &now,
TempUnschedulableUntil: &now,
Extra: map[string]any{
"model_rate_limits": map[string]any{
"claude-sonnet-4-5": map[string]any{
"rate_limit_reset_at": now.Format(time.RFC3339),
},
},
"antigravity_quota_scopes": map[string]any{"gemini": true},
},
},
}
cache := &tempUnschedCacheRecorder{}
svc := NewRateLimitService(repo, nil, &config.Config{}, nil, cache)
result, err := svc.RecoverAccountAfterSuccessfulTest(context.Background(), 42)
require.NoError(t, err)
require.NotNil(t, result)
require.True(t, result.ClearedError)
require.True(t, result.ClearedRateLimit)
require.Equal(t, 1, repo.getByIDCalls)
require.Equal(t, 1, repo.clearErrorCalls)
require.Equal(t, 1, repo.clearRateLimitCalls)
require.Equal(t, 1, repo.clearAntigravityCalls)
require.Equal(t, 1, repo.clearModelRateLimitCalls)
require.Equal(t, 1, repo.clearTempUnschedCalls)
require.Equal(t, []int64{42}, cache.deletedIDs)
}
func TestRateLimitService_RecoverAccountAfterSuccessfulTest_NoRecoverableStateIsNoop(t *testing.T) {
repo := &rateLimitClearRepoStub{
getByIDAccount: &Account{
ID: 7,
Status: StatusActive,
Schedulable: true,
Extra: map[string]any{},
},
}
cache := &tempUnschedCacheRecorder{}
svc := NewRateLimitService(repo, nil, &config.Config{}, nil, cache)
result, err := svc.RecoverAccountAfterSuccessfulTest(context.Background(), 7)
require.NoError(t, err)
require.NotNil(t, result)
require.False(t, result.ClearedError)
require.False(t, result.ClearedRateLimit)
require.Equal(t, 1, repo.getByIDCalls)
require.Equal(t, 0, repo.clearErrorCalls)
require.Equal(t, 0, repo.clearRateLimitCalls)
require.Equal(t, 0, repo.clearAntigravityCalls)
require.Equal(t, 0, repo.clearModelRateLimitCalls)
require.Equal(t, 0, repo.clearTempUnschedCalls)
require.Empty(t, cache.deletedIDs)
}
func TestRateLimitService_RecoverAccountAfterSuccessfulTest_ClearErrorFailed(t *testing.T) {
repo := &rateLimitClearRepoStub{
getByIDAccount: &Account{
ID: 9,
Status: StatusError,
},
clearErrorErr: errors.New("clear error failed"),
}
svc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
result, err := svc.RecoverAccountAfterSuccessfulTest(context.Background(), 9)
require.Error(t, err)
require.Nil(t, result)
require.Equal(t, 1, repo.getByIDCalls)
require.Equal(t, 1, repo.clearErrorCalls)
require.Equal(t, 0, repo.clearRateLimitCalls)
}
func TestRateLimitService_RecoverAccountState_InvalidatesOAuthTokenOnErrorRecovery(t *testing.T) {
repo := &rateLimitClearRepoStub{
getByIDAccount: &Account{
ID: 21,
Type: AccountTypeOAuth,
Status: StatusError,
},
}
invalidator := &recoverTokenInvalidatorStub{}
svc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
svc.SetTokenCacheInvalidator(invalidator)
result, err := svc.RecoverAccountState(context.Background(), 21, AccountRecoveryOptions{
InvalidateToken: true,
})
require.NoError(t, err)
require.NotNil(t, result)
require.True(t, result.ClearedError)
require.False(t, result.ClearedRateLimit)
require.Equal(t, 1, repo.clearErrorCalls)
require.Len(t, invalidator.accounts, 1)
require.Equal(t, int64(21), invalidator.accounts[0].ID)
}

1
backend/internal/service/scheduled_test_port.go
View File

@@ -13,6 +13,7 @@ type ScheduledTestPlan struct {
CronExpression string `json:"cron_expression"`
Enabled bool `json:"enabled"`
MaxResults int `json:"max_results"`
AutoRecover bool `json:"auto_recover"`
LastRunAt *time.Time `json:"last_run_at"`
NextRunAt *time.Time `json:"next_run_at"`
CreatedAt time.Time `json:"created_at"`

31
backend/internal/service/scheduled_test_runner_service.go
View File

@@ -17,6 +17,7 @@ type ScheduledTestRunnerService struct {
planRepo ScheduledTestPlanRepository
scheduledSvc *ScheduledTestService
accountTestSvc *AccountTestService
rateLimitSvc *RateLimitService
cfg *config.Config
cron *cron.Cron
@@ -29,12 +30,14 @@ func NewScheduledTestRunnerService(
planRepo ScheduledTestPlanRepository,
scheduledSvc *ScheduledTestService,
accountTestSvc *AccountTestService,
rateLimitSvc *RateLimitService,
cfg *config.Config,
) *ScheduledTestRunnerService {
return &ScheduledTestRunnerService{
planRepo: planRepo,
scheduledSvc: scheduledSvc,
accountTestSvc: accountTestSvc,
rateLimitSvc: rateLimitSvc,
cfg: cfg,
}
}
@@ -127,6 +130,11 @@ func (s *ScheduledTestRunnerService) runOnePlan(ctx context.Context, plan *Sched
logger.LegacyPrintf("service.scheduled_test_runner", "[ScheduledTestRunner] plan=%d SaveResult error: %v", plan.ID, err)
}
// Auto-recover account if test succeeded and auto_recover is enabled.
if result.Status == "success" && plan.AutoRecover {
s.tryRecoverAccount(ctx, plan.AccountID, plan.ID)
}
nextRun, err := computeNextRun(plan.CronExpression, time.Now())
if err != nil {
logger.LegacyPrintf("service.scheduled_test_runner", "[ScheduledTestRunner] plan=%d computeNextRun error: %v", plan.ID, err)
@@ -137,3 +145,26 @@ func (s *ScheduledTestRunnerService) runOnePlan(ctx context.Context, plan *Sched
logger.LegacyPrintf("service.scheduled_test_runner", "[ScheduledTestRunner] plan=%d UpdateAfterRun error: %v", plan.ID, err)
}
}
// tryRecoverAccount attempts to recover an account from recoverable runtime state.
func (s *ScheduledTestRunnerService) tryRecoverAccount(ctx context.Context, accountID int64, planID int64) {
if s.rateLimitSvc == nil {
return
}
recovery, err := s.rateLimitSvc.RecoverAccountAfterSuccessfulTest(ctx, accountID)
if err != nil {
logger.LegacyPrintf("service.scheduled_test_runner", "[ScheduledTestRunner] plan=%d auto-recover failed: %v", planID, err)
return
}
if recovery == nil {
return
}
if recovery.ClearedError {
logger.LegacyPrintf("service.scheduled_test_runner", "[ScheduledTestRunner] plan=%d auto-recover: account=%d recovered from error status", planID, accountID)
}
if recovery.ClearedRateLimit {
logger.LegacyPrintf("service.scheduled_test_runner", "[ScheduledTestRunner] plan=%d auto-recover: account=%d cleared rate-limit/runtime state", planID, accountID)
}
}

3
backend/internal/service/wire.go
View File

@@ -287,9 +287,10 @@ func ProvideScheduledTestRunnerService(
planRepo ScheduledTestPlanRepository,
scheduledSvc *ScheduledTestService,
accountTestSvc *AccountTestService,
rateLimitSvc *RateLimitService,
cfg *config.Config,
) *ScheduledTestRunnerService {
svc := NewScheduledTestRunnerService(planRepo, scheduledSvc, accountTestSvc, cfg)
svc := NewScheduledTestRunnerService(planRepo, scheduledSvc, accountTestSvc, rateLimitSvc, cfg)
svc.Start()
return svc
}