opensource/sub2api
3
0
Fork 0
mirror of https://github.com/Wei-Shaw/sub2api.git synced 2026-09-24 16:05:44 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #3094 from okbexx/fix-openai-sticky-group-validation

Browse Source 
fix: ignore stale OpenAI sticky sessions outside request group
This commit is contained in:
Wesley Liddick
2026-06-08 08:49:09 +08:00
committed by GitHub
parent 2a7c0fbb0f a436296357
commit 7b394ed157
4 changed files with 224 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/internal/service/openai_account_scheduler.go
+21 -1
View File
@@ -376,7 +376,7 @@ func (s *defaultOpenAIAccountScheduler) selectBySessionHash(
return nil, false, nil
}
account = s.service.recheckSelectedOpenAIAccountFromDB(ctx, account, req.RequestedModel, req.RequireCompact, req.RequiredCapability)
if account == nil || !s.isAccountTransportCompatible(account, req.RequiredTransport) {
if account == nil || !openAIStickyAccountMatchesGroup(account, req.GroupID) || !s.isAccountTransportCompatible(account, req.RequiredTransport) {
_ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash)
return nil, false, nil
}
@@ -426,6 +426,26 @@ func (s *defaultOpenAIAccountScheduler) selectBySessionHash(
return nil, false, nil
}
func openAIStickyAccountMatchesGroup(account *Account, groupID *int64) bool {
if account == nil {
return false
}
if groupID == nil {
return len(account.AccountGroups) == 0 && len(account.GroupIDs) == 0
}
for _, accountGroupID := range account.GroupIDs {
if accountGroupID == *groupID {
return true
}
}
for _, accountGroup := range account.AccountGroups {
if accountGroup.GroupID == *groupID {
return true
}
}
return false
}
func (s *defaultOpenAIAccountScheduler) shouldEscapeStickyAccount(accountID int64, cfg openAIStickyEscapeConfig) (reason string, errorRate float64, ttft float64, shouldEscape bool) {
if !cfg.enabled || s == nil || s.stats == nil || accountID <= 0 {
return "", 0, 0, false
backend/internal/service/openai_account_scheduler_test.go
+102 -6
View File
@@ -57,6 +57,30 @@ func (r schedulerTestOpenAIAccountRepo) ListSchedulableUngroupedByPlatform(ctx c
return r.ListSchedulableByPlatform(ctx, platform)
}
type schedulerGroupAwareOpenAIAccountRepo struct {
schedulerTestOpenAIAccountRepo
}
func (r schedulerGroupAwareOpenAIAccountRepo) ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]Account, error) {
var result []Account
for _, acc := range r.accounts {
if acc.Platform == platform && openAIStickyAccountMatchesGroup(&acc, &groupID) {
result = append(result, acc)
}
}
return result, nil
}
func (r schedulerGroupAwareOpenAIAccountRepo) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]Account, error) {
var result []Account
for _, acc := range r.accounts {
if acc.Platform == platform && openAIStickyAccountMatchesGroup(&acc, nil) {
result = append(result, acc)
}
}
return result, nil
}
type schedulerTestConcurrencyCache struct {
ConcurrencyCache
loadBatchErr error
@@ -1162,6 +1186,7 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionSticky(t *testin
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
GroupIDs: []int64{groupID},
}
cache := &schedulerTestGatewayCache{
sessionBindings: map[string]int64{
@@ -1210,6 +1235,7 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyBusyKeepsS
Schedulable: true,
Concurrency: 1,
Priority: 0,
GroupIDs: []int64{groupID},
},
{
ID: 21002,
@@ -1219,6 +1245,7 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyBusyKeepsS
Schedulable: true,
Concurrency: 1,
Priority: 9,
GroupIDs: []int64{groupID},
},
}
cache := &schedulerTestGatewayCache{
@@ -1292,6 +1319,7 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyEscapeByTT
Schedulable: true,
Concurrency: 1,
Priority: 0,
GroupIDs: []int64{groupID},
},
{
ID: 21102,
@@ -1301,6 +1329,7 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyEscapeByTT
Schedulable: true,
Concurrency: 1,
Priority: 1,
GroupIDs: []int64{groupID},
},
}
cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{"openai:session_hash_sticky_ttft": 21101}}
@@ -1355,8 +1384,8 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyEscapeByEr
ctx := context.Background()
groupID := int64(10102)
accounts := []Account{
{ID: 21201, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0},
{ID: 21202, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
{ID: 21201, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}},
{ID: 21202, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1, GroupIDs: []int64{groupID}},
}
cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{"openai:session_hash_sticky_error_rate": 21201}}
cfg := &config.Config{}
@@ -1405,8 +1434,8 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyBusyEscape
ctx := context.Background()
groupID := int64(10103)
accounts := []Account{
{ID: 21301, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0},
{ID: 21302, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
{ID: 21301, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}},
{ID: 21302, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1, GroupIDs: []int64{groupID}},
}
cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{"openai:session_hash_sticky_busy_escape": 21301}}
cfg := &config.Config{}
@@ -1448,8 +1477,8 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyEscapeDisa
ctx := context.Background()
groupID := int64(10104)
accounts := []Account{
{ID: 21401, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0},
{ID: 21402, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
{ID: 21401, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}},
{ID: 21402, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1, GroupIDs: []int64{groupID}},
}
cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{"openai:session_hash_sticky_disabled": 21401}}
cfg := &config.Config{}
@@ -1537,6 +1566,7 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionSticky_ForceHTTP
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
GroupIDs: []int64{groupID},
Extra: map[string]any{
"openai_ws_force_http": true,
},
@@ -1588,6 +1618,7 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_RequiredWSV2_SkipsStick
Schedulable: true,
Concurrency: 1,
Priority: 0,
GroupIDs: []int64{groupID},
},
{
ID: 2202,
@@ -1597,6 +1628,7 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_RequiredWSV2_SkipsStick
Schedulable: true,
Concurrency: 1,
Priority: 5,
GroupIDs: []int64{groupID},
Extra: map[string]any{
"openai_apikey_responses_websockets_v2_enabled": true,
},
@@ -1647,6 +1679,69 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_RequiredWSV2_SkipsStick
}
}
func TestOpenAIGatewayService_SelectAccountWithScheduler_ClearsStickyAccountOutsideGroup(t *testing.T) {
ctx := context.Background()
groupID := int64(1013)
accounts := []Account{
{
ID: 2401,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Priority: 0,
},
{
ID: 2402,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Priority: 5,
AccountGroups: []AccountGroup{
{AccountID: 2402, GroupID: groupID},
},
},
}
cache := &schedulerTestGatewayCache{
sessionBindings: map[string]int64{
"openai:session_hash_removed_group": 2401,
},
}
svc := &OpenAIGatewayService{
accountRepo: schedulerGroupAwareOpenAIAccountRepo{schedulerTestOpenAIAccountRepo{accounts: accounts}},
cache: cache,
cfg: &config.Config{},
rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
}
selection, decision, err := svc.SelectAccountWithScheduler(
ctx,
&groupID,
"",
"session_hash_removed_group",
"gpt-5.1",
nil,
OpenAIUpstreamTransportAny,
false,
)
require.NoError(t, err)
require.NotNil(t, selection)
require.NotNil(t, selection.Account)
require.Equal(t, int64(2402), selection.Account.ID)
require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
require.False(t, decision.StickySessionHit)
require.Equal(t, 1, cache.deletedSessions["openai:session_hash_removed_group"])
require.Equal(t, int64(2402), cache.sessionBindings["openai:session_hash_removed_group"])
if selection.ReleaseFunc != nil {
selection.ReleaseFunc()
}
}
func TestOpenAIGatewayService_SelectAccountWithScheduler_RequiredWSV2_NoAvailableAccount(t *testing.T) {
ctx := context.Background()
groupID := int64(1012)
@@ -1858,6 +1953,7 @@ func TestOpenAIGatewayService_OpenAIAccountSchedulerMetrics(t *testing.T) {
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
GroupIDs: []int64{groupID},
}
cache := &schedulerTestGatewayCache{
sessionBindings: map[string]int64{
backend/internal/service/openai_gateway_service.go
+3 -1
View File
@@ -1712,7 +1712,7 @@ func (s *OpenAIGatewayService) tryStickySessionHit(ctx context.Context, groupID
return nil
}
account = s.recheckSelectedOpenAIAccountFromDB(ctx, account, requestedModel, requireCompact, requiredCapability)
if account == nil {
if account == nil || !openAIStickyAccountMatchesGroup(account, groupID) {
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
return nil
}
@@ -1908,6 +1908,8 @@ func (s *OpenAIGatewayService) selectAccountWithLoadAwareness(ctx context.Contex
account = s.recheckSelectedOpenAIAccountFromDB(ctx, account, requestedModel, requireCompact, requiredCapability)
if account == nil {
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
} else if !openAIStickyAccountMatchesGroup(account, groupID) {
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
} else if s.isOpenAIAccountRuntimeBlocked(account) {
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
} else if needsUpstreamCheck && s.isUpstreamModelRestrictedByChannel(ctx, *groupID, account, requestedModel, requireCompact) {
backend/internal/service/openai_gateway_service_test.go
+98
View File
@@ -79,6 +79,30 @@ func (r stubOpenAIAccountRepo) ListSchedulableUngroupedByPlatform(ctx context.Co
return r.ListSchedulableByPlatform(ctx, platform)
}
type groupAwareStubOpenAIAccountRepo struct {
stubOpenAIAccountRepo
}
func (r groupAwareStubOpenAIAccountRepo) ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]Account, error) {
var result []Account
for _, acc := range r.accounts {
if acc.Platform == platform && openAIStickyAccountMatchesGroup(&acc, &groupID) {
result = append(result, acc)
}
}
return result, nil
}
func (r groupAwareStubOpenAIAccountRepo) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]Account, error) {
var result []Account
for _, acc := range r.accounts {
if acc.Platform == platform && openAIStickyAccountMatchesGroup(&acc, nil) {
result = append(result, acc)
}
}
return result, nil
}
type stubConcurrencyCache struct {
ConcurrencyCache
loadBatchErr error
@@ -576,6 +600,41 @@ func TestOpenAISelectAccountForModelWithExclusions_StickyUnschedulableClearsSess
}
}
func TestOpenAISelectAccountForModelWithExclusions_StickyOutsideGroupClearsSession(t *testing.T) {
sessionHash := "session-outside-group"
groupID := int64(1001)
repo := groupAwareStubOpenAIAccountRepo{
stubOpenAIAccountRepo{
accounts: []Account{
{ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1},
{ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, AccountGroups: []AccountGroup{{GroupID: groupID}}},
},
},
}
cache := &stubGatewayCache{
sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
}
svc := &OpenAIGatewayService{
accountRepo: repo,
cache: cache,
}
acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), &groupID, sessionHash, "gpt-4", nil)
if err != nil {
t.Fatalf("SelectAccountForModelWithExclusions error: %v", err)
}
if acc == nil || acc.ID != 2 {
t.Fatalf("expected account 2, got %+v", acc)
}
if cache.deletedSessions["openai:"+sessionHash] != 1 {
t.Fatalf("expected sticky session to be deleted")
}
if cache.sessionBindings["openai:"+sessionHash] != 2 {
t.Fatalf("expected sticky session to bind to account 2")
}
}
func TestOpenAISelectAccountWithLoadAwareness_StickyUnschedulableClearsSession(t *testing.T) {
sessionHash := "session-2"
groupID := int64(1)
@@ -613,6 +672,45 @@ func TestOpenAISelectAccountWithLoadAwareness_StickyUnschedulableClearsSession(t
}
}
func TestOpenAISelectAccountWithLoadAwareness_StickyOutsideGroupClearsSession(t *testing.T) {
sessionHash := "session-load-outside-group"
groupID := int64(1002)
repo := groupAwareStubOpenAIAccountRepo{
stubOpenAIAccountRepo{
accounts: []Account{
{ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1},
{ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, AccountGroups: []AccountGroup{{GroupID: groupID}}},
},
},
}
cache := &stubGatewayCache{
sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
}
svc := &OpenAIGatewayService{
accountRepo: repo,
cache: cache,
concurrencyService: NewConcurrencyService(stubConcurrencyCache{}),
}
selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, sessionHash, "gpt-4", nil)
if err != nil {
t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
}
if selection == nil || selection.Account == nil || selection.Account.ID != 2 {
t.Fatalf("expected account 2, got %+v", selection)
}
if cache.deletedSessions["openai:"+sessionHash] != 1 {
t.Fatalf("expected sticky session to be deleted")
}
if cache.sessionBindings["openai:"+sessionHash] != 2 {
t.Fatalf("expected sticky session to bind to account 2")
}
if selection.ReleaseFunc != nil {
selection.ReleaseFunc()
}
}
func TestOpenAISelectAccountForModelWithExclusions_NoModelSupport(t *testing.T) {
repo := stubOpenAIAccountRepo{
accounts: []Account{