Skip to content

Commit c7137df

Browse files
Wei-ShawWei-Shaw
authored
Merge pull request Wei-Shaw#1218 from LvyuanW/openai-runtime-recheck
fix(openai): prevent rescheduling rate-limited accounts
2 parents 5a3375c + fef9259 commit c7137df

17 files changed

Lines changed: 372 additions & 57 deletions

‎backend/internal/repository/account_repo.go‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -404,6 +404,17 @@ func (r *accountRepository) Update(ctx context.Context, account *service.Account
404404
return nil
405405
}
406406

407+
func (r *accountRepository) UpdateCredentials(ctx context.Context, id int64, credentials map[string]any) error {
408+
_, err := r.client.Account.UpdateOneID(id).
409+
SetCredentials(normalizeJSONMap(credentials)).
410+
Save(ctx)
411+
if err != nil {
412+
return translatePersistenceError(err, service.ErrAccountNotFound, nil)
413+
}
414+
r.syncSchedulerAccountSnapshot(ctx, id)
415+
return nil
416+
}
417+
407418
func (r *accountRepository) Delete(ctx context.Context, id int64) error {
408419
groupIDs, err := r.loadAccountGroupIDs(ctx, id)
409420
if err != nil {

‎backend/internal/service/account_credentials_persistence.go‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
package service
2+
3+
import "context"
4+
5+
type accountCredentialsUpdater interface {
6+
UpdateCredentials(ctx context.Context, id int64, credentials map[string]any) error
7+
}
8+
9+
func persistAccountCredentials(ctx context.Context, repo AccountRepository, account *Account, credentials map[string]any) error {
10+
if repo == nil || account == nil {
11+
return nil
12+
}
13+
14+
account.Credentials = cloneCredentials(credentials)
15+
if updater, ok := any(repo).(accountCredentialsUpdater); ok {
16+
return updater.UpdateCredentials(ctx, account.ID, account.Credentials)
17+
}
18+
return repo.Update(ctx, account)
19+
}
20+
21+
func cloneCredentials(in map[string]any) map[string]any {
22+
if in == nil {
23+
return map[string]any{}
24+
}
25+
out := make(map[string]any, len(in))
26+
for k, v := range in {
27+
out[k] = v
28+
}
29+
return out
30+
}

‎backend/internal/service/antigravity_token_provider.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,7 @@ func (p *AntigravityTokenProvider) GetAccessToken(ctx context.Context, account *
138138
p.markBackfillAttempted(account.ID)
139139
if projectID, err := p.antigravityOAuthService.FillProjectID(ctx, account, accessToken); err == nil && projectID != "" {
140140
account.Credentials["project_id"] = projectID
141-
if updateErr := p.accountRepo.Update(ctx, account); updateErr != nil {
141+
if updateErr := persistAccountCredentials(ctx, p.accountRepo, account, account.Credentials); updateErr != nil {
142142
slog.Warn("antigravity_project_id_backfill_persist_failed",
143143
"account_id", account.ID,
144144
"error", updateErr,

‎backend/internal/service/crs_sync_service.go‎

Lines changed: 6 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -367,8 +367,7 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
367367
// 🔄 Refresh OAuth token after creation
368368
if targetType == AccountTypeOAuth {
369369
if refreshedCreds := s.refreshOAuthToken(ctx, account); refreshedCreds != nil {
370-
account.Credentials = refreshedCreds
371-
_ = s.accountRepo.Update(ctx, account)
370+
_ = persistAccountCredentials(ctx, s.accountRepo, account, refreshedCreds)
372371
}
373372
}
374373
item.Action = "created"
@@ -402,8 +401,7 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
402401
// 🔄 Refresh OAuth token after update
403402
if targetType == AccountTypeOAuth {
404403
if refreshedCreds := s.refreshOAuthToken(ctx, existing); refreshedCreds != nil {
405-
existing.Credentials = refreshedCreds
406-
_ = s.accountRepo.Update(ctx, existing)
404+
_ = persistAccountCredentials(ctx, s.accountRepo, existing, refreshedCreds)
407405
}
408406
}
409407

@@ -620,8 +618,7 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
620618
}
621619
// 🔄 Refresh OAuth token after creation
622620
if refreshedCreds := s.refreshOAuthToken(ctx, account); refreshedCreds != nil {
623-
account.Credentials = refreshedCreds
624-
_ = s.accountRepo.Update(ctx, account)
621+
_ = persistAccountCredentials(ctx, s.accountRepo, account, refreshedCreds)
625622
}
626623
item.Action = "created"
627624
result.Created++
@@ -652,8 +649,7 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
652649

653650
// 🔄 Refresh OAuth token after update
654651
if refreshedCreds := s.refreshOAuthToken(ctx, existing); refreshedCreds != nil {
655-
existing.Credentials = refreshedCreds
656-
_ = s.accountRepo.Update(ctx, existing)
652+
_ = persistAccountCredentials(ctx, s.accountRepo, existing, refreshedCreds)
657653
}
658654

659655
item.Action = "updated"
@@ -862,8 +858,7 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
862858
continue
863859
}
864860
if refreshedCreds := s.refreshOAuthToken(ctx, account); refreshedCreds != nil {
865-
account.Credentials = refreshedCreds
866-
_ = s.accountRepo.Update(ctx, account)
861+
_ = persistAccountCredentials(ctx, s.accountRepo, account, refreshedCreds)
867862
}
868863
item.Action = "created"
869864
result.Created++
@@ -893,8 +888,7 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
893888
}
894889

895890
if refreshedCreds := s.refreshOAuthToken(ctx, existing); refreshedCreds != nil {
896-
existing.Credentials = refreshedCreds
897-
_ = s.accountRepo.Update(ctx, existing)
891+
_ = persistAccountCredentials(ctx, s.accountRepo, existing, refreshedCreds)
898892
}
899893

900894
item.Action = "updated"

‎backend/internal/service/gemini_token_provider.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,7 @@ func (p *GeminiTokenProvider) GetAccessToken(ctx context.Context, account *Accou
135135
if tierID != "" {
136136
account.Credentials["tier_id"] = tierID
137137
}
138-
_ = p.accountRepo.Update(ctx, account)
138+
_ = persistAccountCredentials(ctx, p.accountRepo, account, account.Credentials)
139139
}
140140
}
141141

‎backend/internal/service/oauth_refresh_api.go‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -108,8 +108,7 @@ func (api *OAuthRefreshAPI) RefreshIfNeeded(
108108
// 5. 设置版本号 + 更新 DB
109109
if newCredentials != nil {
110110
newCredentials["_token_version"] = time.Now().UnixMilli()
111-
freshAccount.Credentials = newCredentials
112-
if updateErr := api.accountRepo.Update(ctx, freshAccount); updateErr != nil {
111+
if updateErr := persistAccountCredentials(ctx, api.accountRepo, freshAccount, newCredentials); updateErr != nil {
113112
slog.Error("oauth_refresh_update_failed",
114113
"account_id", freshAccount.ID,
115114
"error", updateErr,

‎backend/internal/service/oauth_refresh_api_test.go‎

Lines changed: 48 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,11 @@ import (
1616
// refreshAPIAccountRepo implements AccountRepository for OAuthRefreshAPI tests.
1717
type refreshAPIAccountRepo struct {
1818
mockAccountRepoForGemini
19-
account *Account // returned by GetByID
20-
getByIDErr error
21-
updateErr error
22-
updateCalls int
19+
account *Account // returned by GetByID
20+
getByIDErr error
21+
updateErr error
22+
updateCalls int
23+
updateCredentialsCalls int
2324
}
2425

2526
func (r *refreshAPIAccountRepo) GetByID(_ context.Context, _ int64) (*Account, error) {
@@ -34,6 +35,19 @@ func (r *refreshAPIAccountRepo) Update(_ context.Context, _ *Account) error {
3435
return r.updateErr
3536
}
3637

38+
func (r *refreshAPIAccountRepo) UpdateCredentials(_ context.Context, id int64, credentials map[string]any) error {
39+
r.updateCalls++
40+
r.updateCredentialsCalls++
41+
if r.updateErr != nil {
42+
return r.updateErr
43+
}
44+
if r.account == nil || r.account.ID != id {
45+
r.account = &Account{ID: id}
46+
}
47+
r.account.Credentials = cloneCredentials(credentials)
48+
return nil
49+
}
50+
3751
// refreshAPIExecutorStub implements OAuthRefreshExecutor for tests.
3852
type refreshAPIExecutorStub struct {
3953
needsRefresh bool
@@ -106,10 +120,36 @@ func TestRefreshIfNeeded_Success(t *testing.T) {
106120
require.Equal(t, "new-token", result.NewCredentials["access_token"])
107121
require.NotNil(t, result.NewCredentials["_token_version"]) // version stamp set
108122
require.Equal(t, 1, repo.updateCalls) // DB updated
109-
require.Equal(t, 1, cache.releaseCalls) // lock released
123+
require.Equal(t, 1, repo.updateCredentialsCalls)
124+
require.Equal(t, 1, cache.releaseCalls) // lock released
110125
require.Equal(t, 1, executor.refreshCalls)
111126
}
112127

128+
func TestRefreshIfNeeded_UpdateCredentialsPreservesRateLimitState(t *testing.T) {
129+
resetAt := time.Now().Add(45 * time.Minute)
130+
account := &Account{
131+
ID: 11,
132+
Platform: PlatformGemini,
133+
Type: AccountTypeOAuth,
134+
RateLimitResetAt: &resetAt,
135+
}
136+
repo := &refreshAPIAccountRepo{account: account}
137+
cache := &refreshAPICacheStub{lockResult: true}
138+
executor := &refreshAPIExecutorStub{
139+
needsRefresh: true,
140+
credentials: map[string]any{"access_token": "safe-token"},
141+
}
142+
143+
api := NewOAuthRefreshAPI(repo, cache)
144+
result, err := api.RefreshIfNeeded(context.Background(), account, executor, 3*time.Minute)
145+
146+
require.NoError(t, err)
147+
require.True(t, result.Refreshed)
148+
require.Equal(t, 1, repo.updateCredentialsCalls)
149+
require.NotNil(t, repo.account.RateLimitResetAt)
150+
require.WithinDuration(t, resetAt, *repo.account.RateLimitResetAt, time.Second)
151+
}
152+
113153
func TestRefreshIfNeeded_LockHeld(t *testing.T) {
114154
account := &Account{ID: 2, Platform: PlatformAnthropic}
115155
repo := &refreshAPIAccountRepo{account: account}
@@ -193,7 +233,7 @@ func TestRefreshIfNeeded_RefreshError(t *testing.T) {
193233
require.Error(t, err)
194234
require.Nil(t, result)
195235
require.Contains(t, err.Error(), "invalid_grant")
196-
require.Equal(t, 0, repo.updateCalls) // no DB update on refresh error
236+
require.Equal(t, 0, repo.updateCalls) // no DB update on refresh error
197237
require.Equal(t, 1, cache.releaseCalls) // lock still released via defer
198238
}
199239

@@ -299,8 +339,8 @@ func TestMergeCredentials_NewOverridesOld(t *testing.T) {
299339

300340
result := MergeCredentials(old, new)
301341

302-
require.Equal(t, "new-token", result["access_token"]) // overridden
303-
require.Equal(t, "old-refresh", result["refresh_token"]) // preserved
342+
require.Equal(t, "new-token", result["access_token"]) // overridden
343+
require.Equal(t, "old-refresh", result["refresh_token"]) // preserved
304344
}
305345

306346
// ========== BuildClaudeAccountCredentials tests ==========

‎backend/internal/service/openai_account_scheduler.go‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -330,6 +330,11 @@ func (s *defaultOpenAIAccountScheduler) selectBySessionHash(
330330
_ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash)
331331
return nil, nil
332332
}
333+
account = s.service.recheckSelectedOpenAIAccountFromDB(ctx, account, req.RequestedModel)
334+
if account == nil {
335+
_ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash)
336+
return nil, nil
337+
}
333338

334339
result, acquireErr := s.service.tryAcquireAccountSlot(ctx, accountID, account.Concurrency)
335340
if acquireErr == nil && result.Acquired {
@@ -691,6 +696,10 @@ func (s *defaultOpenAIAccountScheduler) selectByLoadBalance(
691696
if fresh == nil || !s.isAccountTransportCompatible(fresh, req.RequiredTransport) {
692697
continue
693698
}
699+
fresh = s.service.recheckSelectedOpenAIAccountFromDB(ctx, fresh, req.RequestedModel)
700+
if fresh == nil || !s.isAccountTransportCompatible(fresh, req.RequiredTransport) {
701+
continue
702+
}
694703
result, acquireErr := s.service.tryAcquireAccountSlot(ctx, fresh.ID, fresh.Concurrency)
695704
if acquireErr != nil {
696705
return nil, len(candidates), topK, loadSkew, acquireErr

‎backend/internal/service/openai_account_scheduler_test.go‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,61 @@ func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_SkipsFreshlyRa
8484
require.Equal(t, int64(32002), account.ID)
8585
}
8686

87+
func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyDBRuntimeRecheckSkipsStaleCachedAccount(t *testing.T) {
88+
ctx := context.Background()
89+
groupID := int64(10103)
90+
rateLimitedUntil := time.Now().Add(30 * time.Minute)
91+
staleSticky := &Account{ID: 33001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0}
92+
staleBackup := &Account{ID: 33002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
93+
dbSticky := Account{ID: 33001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, RateLimitResetAt: &rateLimitedUntil}
94+
dbBackup := Account{ID: 33002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
95+
cache := &stubGatewayCache{sessionBindings: map[string]int64{"openai:session_hash_db_runtime_recheck": 33001}}
96+
snapshotCache := &openAISnapshotCacheStub{
97+
snapshotAccounts: []*Account{staleSticky, staleBackup},
98+
accountsByID: map[int64]*Account{33001: staleSticky, 33002: staleBackup},
99+
}
100+
snapshotService := &SchedulerSnapshotService{cache: snapshotCache}
101+
svc := &OpenAIGatewayService{
102+
accountRepo: stubOpenAIAccountRepo{accounts: []Account{dbSticky, dbBackup}},
103+
cache: cache,
104+
cfg: &config.Config{},
105+
schedulerSnapshot: snapshotService,
106+
concurrencyService: NewConcurrencyService(stubConcurrencyCache{}),
107+
}
108+
109+
selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_hash_db_runtime_recheck", "gpt-5.1", nil, OpenAIUpstreamTransportAny)
110+
require.NoError(t, err)
111+
require.NotNil(t, selection)
112+
require.NotNil(t, selection.Account)
113+
require.Equal(t, int64(33002), selection.Account.ID)
114+
require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
115+
}
116+
117+
func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_DBRuntimeRecheckSkipsStaleCachedCandidate(t *testing.T) {
118+
ctx := context.Background()
119+
groupID := int64(10104)
120+
rateLimitedUntil := time.Now().Add(30 * time.Minute)
121+
stalePrimary := &Account{ID: 34001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0}
122+
staleSecondary := &Account{ID: 34002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
123+
dbPrimary := Account{ID: 34001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, RateLimitResetAt: &rateLimitedUntil}
124+
dbSecondary := Account{ID: 34002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
125+
snapshotCache := &openAISnapshotCacheStub{
126+
snapshotAccounts: []*Account{stalePrimary, staleSecondary},
127+
accountsByID: map[int64]*Account{34001: stalePrimary, 34002: staleSecondary},
128+
}
129+
snapshotService := &SchedulerSnapshotService{cache: snapshotCache}
130+
svc := &OpenAIGatewayService{
131+
accountRepo: stubOpenAIAccountRepo{accounts: []Account{dbPrimary, dbSecondary}},
132+
cfg: &config.Config{},
133+
schedulerSnapshot: snapshotService,
134+
}
135+
136+
account, err := svc.SelectAccountForModelWithExclusions(ctx, &groupID, "", "gpt-5.1", nil)
137+
require.NoError(t, err)
138+
require.NotNil(t, account)
139+
require.Equal(t, int64(34002), account.ID)
140+
}
141+
87142
func TestOpenAIGatewayService_SelectAccountWithScheduler_PreviousResponseSticky(t *testing.T) {
88143
ctx := context.Background()
89144
groupID := int64(9)

0 commit comments

Comments
 (0)