Skip to content

Commit fef9259

Browse files
LvyuanWLvyuanW
committed
fix(openai): recheck runtime state from db before final account selection
1 parent ad7c107 commit fef9259

5 files changed

Lines changed: 177 additions & 20 deletions

File tree

‎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)

‎backend/internal/service/openai_gateway_service.go‎

Lines changed: 56 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1201,6 +1201,11 @@ func (s *OpenAIGatewayService) tryStickySessionHit(ctx context.Context, groupID
12011201
if requestedModel != "" && !account.IsModelSupported(requestedModel) {
12021202
return nil
12031203
}
1204+
account = s.recheckSelectedOpenAIAccountFromDB(ctx, account, requestedModel)
1205+
if account == nil {
1206+
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
1207+
return nil
1208+
}
12041209

12051210
// 刷新会话 TTL 并返回账号
12061211
// Refresh session TTL and return account
@@ -1229,6 +1234,10 @@ func (s *OpenAIGatewayService) selectBestAccount(ctx context.Context, accounts [
12291234
if fresh == nil {
12301235
continue
12311236
}
1237+
fresh = s.recheckSelectedOpenAIAccountFromDB(ctx, fresh, requestedModel)
1238+
if fresh == nil {
1239+
continue
1240+
}
12321241

12331242
// 选择优先级最高且最久未使用的账号
12341243
// Select highest priority and least recently used
@@ -1353,27 +1362,32 @@ func (s *OpenAIGatewayService) SelectAccountWithLoadAwareness(ctx context.Contex
13531362
}
13541363
if !clearSticky && account.IsSchedulable() && account.IsOpenAI() &&
13551364
(requestedModel == "" || account.IsModelSupported(requestedModel)) {
1356-
result, err := s.tryAcquireAccountSlot(ctx, accountID, account.Concurrency)
1357-
if err == nil && result.Acquired {
1358-
_ = s.refreshStickySessionTTL(ctx, groupID, sessionHash, openaiStickySessionTTL)
1359-
return &AccountSelectionResult{
1360-
Account: account,
1361-
Acquired: true,
1362-
ReleaseFunc: result.ReleaseFunc,
1363-
}, nil
1364-
}
1365+
account = s.recheckSelectedOpenAIAccountFromDB(ctx, account, requestedModel)
1366+
if account == nil {
1367+
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
1368+
} else {
1369+
result, err := s.tryAcquireAccountSlot(ctx, accountID, account.Concurrency)
1370+
if err == nil && result.Acquired {
1371+
_ = s.refreshStickySessionTTL(ctx, groupID, sessionHash, openaiStickySessionTTL)
1372+
return &AccountSelectionResult{
1373+
Account: account,
1374+
Acquired: true,
1375+
ReleaseFunc: result.ReleaseFunc,
1376+
}, nil
1377+
}
13651378

1366-
waitingCount, _ := s.concurrencyService.GetAccountWaitingCount(ctx, accountID)
1367-
if waitingCount < cfg.StickySessionMaxWaiting {
1368-
return &AccountSelectionResult{
1369-
Account: account,
1370-
WaitPlan: &AccountWaitPlan{
1371-
AccountID: accountID,
1372-
MaxConcurrency: account.Concurrency,
1373-
Timeout: cfg.StickySessionWaitTimeout,
1374-
MaxWaiting: cfg.StickySessionMaxWaiting,
1375-
},
1376-
}, nil
1379+
waitingCount, _ := s.concurrencyService.GetAccountWaitingCount(ctx, accountID)
1380+
if waitingCount < cfg.StickySessionMaxWaiting {
1381+
return &AccountSelectionResult{
1382+
Account: account,
1383+
WaitPlan: &AccountWaitPlan{
1384+
AccountID: accountID,
1385+
MaxConcurrency: account.Concurrency,
1386+
Timeout: cfg.StickySessionWaitTimeout,
1387+
MaxWaiting: cfg.StickySessionMaxWaiting,
1388+
},
1389+
}, nil
1390+
}
13771391
}
13781392
}
13791393
}
@@ -1560,6 +1574,28 @@ func (s *OpenAIGatewayService) resolveFreshSchedulableOpenAIAccount(ctx context.
15601574
return fresh
15611575
}
15621576

1577+
func (s *OpenAIGatewayService) recheckSelectedOpenAIAccountFromDB(ctx context.Context, account *Account, requestedModel string) *Account {
1578+
if account == nil {
1579+
return nil
1580+
}
1581+
if s.schedulerSnapshot == nil || s.accountRepo == nil {
1582+
return account
1583+
}
1584+
1585+
latest, err := s.accountRepo.GetByID(ctx, account.ID)
1586+
if err != nil || latest == nil {
1587+
return nil
1588+
}
1589+
syncOpenAICodexRateLimitFromExtra(ctx, s.accountRepo, latest, time.Now())
1590+
if !latest.IsSchedulable() || !latest.IsOpenAI() {
1591+
return nil
1592+
}
1593+
if requestedModel != "" && !latest.IsModelSupported(requestedModel) {
1594+
return nil
1595+
}
1596+
return latest
1597+
}
1598+
15631599
func (s *OpenAIGatewayService) getSchedulableAccount(ctx context.Context, accountID int64) (*Account, error) {
15641600
var (
15651601
account *Account

‎backend/internal/service/openai_ws_account_sticky_test.go‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,58 @@ func TestOpenAIGatewayService_SelectAccountByPreviousResponseID_RateLimitedMiss(
8585
require.Zero(t, boundAccountID)
8686
}
8787

88+
func TestOpenAIGatewayService_SelectAccountByPreviousResponseID_DBRuntimeRecheckRateLimitedMiss(t *testing.T) {
89+
ctx := context.Background()
90+
groupID := int64(24)
91+
rateLimitedUntil := time.Now().Add(30 * time.Minute)
92+
staleAccount := &Account{
93+
ID: 13,
94+
Platform: PlatformOpenAI,
95+
Type: AccountTypeAPIKey,
96+
Status: StatusActive,
97+
Schedulable: true,
98+
Concurrency: 1,
99+
Extra: map[string]any{
100+
"openai_apikey_responses_websockets_v2_enabled": true,
101+
},
102+
}
103+
dbAccount := Account{
104+
ID: 13,
105+
Platform: PlatformOpenAI,
106+
Type: AccountTypeAPIKey,
107+
Status: StatusActive,
108+
Schedulable: true,
109+
Concurrency: 1,
110+
RateLimitResetAt: &rateLimitedUntil,
111+
Extra: map[string]any{
112+
"openai_apikey_responses_websockets_v2_enabled": true,
113+
},
114+
}
115+
cache := &stubGatewayCache{}
116+
store := NewOpenAIWSStateStore(cache)
117+
cfg := newOpenAIWSV2TestConfig()
118+
snapshotCache := &openAISnapshotCacheStub{
119+
accountsByID: map[int64]*Account{dbAccount.ID: staleAccount},
120+
}
121+
svc := &OpenAIGatewayService{
122+
accountRepo: stubOpenAIAccountRepo{accounts: []Account{dbAccount}},
123+
cache: cache,
124+
cfg: cfg,
125+
concurrencyService: NewConcurrencyService(stubConcurrencyCache{}),
126+
openaiWSStateStore: store,
127+
schedulerSnapshot: &SchedulerSnapshotService{cache: snapshotCache},
128+
}
129+
130+
require.NoError(t, store.BindResponseAccount(ctx, groupID, "resp_prev_db_rl", dbAccount.ID, time.Hour))
131+
132+
selection, err := svc.SelectAccountByPreviousResponseID(ctx, &groupID, "resp_prev_db_rl", "gpt-5.1", nil)
133+
require.NoError(t, err)
134+
require.Nil(t, selection, "DB 中已限流的账号不应继续命中 previous_response_id 粘连")
135+
boundAccountID, getErr := store.GetResponseAccount(ctx, groupID, "resp_prev_db_rl")
136+
require.NoError(t, getErr)
137+
require.Zero(t, boundAccountID)
138+
}
139+
88140
func TestOpenAIGatewayService_SelectAccountByPreviousResponseID_Excluded(t *testing.T) {
89141
ctx := context.Background()
90142
groupID := int64(23)

‎backend/internal/service/openai_ws_forwarder.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3846,6 +3846,11 @@ func (s *OpenAIGatewayService) SelectAccountByPreviousResponseID(
38463846
if requestedModel != "" && !account.IsModelSupported(requestedModel) {
38473847
return nil, nil
38483848
}
3849+
account = s.recheckSelectedOpenAIAccountFromDB(ctx, account, requestedModel)
3850+
if account == nil {
3851+
_ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID)
3852+
return nil, nil
3853+
}
38493854

38503855
result, acquireErr := s.tryAcquireAccountSlot(ctx, accountID, account.Concurrency)
38513856
if acquireErr == nil && result.Acquired {

0 commit comments

Comments
 (0)