diff --git a/backend/internal/service/openai_account_scheduler.go b/backend/internal/service/openai_account_scheduler.go index 6e1ab3fda6..ece3716b57 100644 --- a/backend/internal/service/openai_account_scheduler.go +++ b/backend/internal/service/openai_account_scheduler.go @@ -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 diff --git a/backend/internal/service/openai_account_scheduler_test.go b/backend/internal/service/openai_account_scheduler_test.go index 845d60c41a..4d2fbd7d70 100644 --- a/backend/internal/service/openai_account_scheduler_test.go +++ b/backend/internal/service/openai_account_scheduler_test.go @@ -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{ diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index 062a54bf40..5139813553 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -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) { diff --git a/backend/internal/service/openai_gateway_service_test.go b/backend/internal/service/openai_gateway_service_test.go index e00a3dc2e2..9866ddeff4 100644 --- a/backend/internal/service/openai_gateway_service_test.go +++ b/backend/internal/service/openai_gateway_service_test.go @@ -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{