diff --git a/backend/internal/service/openai_gateway_scheduling.go b/backend/internal/service/openai_gateway_scheduling.go index 7ab83a3de8..a90bad3e25 100644 --- a/backend/internal/service/openai_gateway_scheduling.go +++ b/backend/internal/service/openai_gateway_scheduling.go @@ -1174,6 +1174,11 @@ func (s *OpenAIGatewayService) selectAccountWithLoadAwareness(ctx context.Contex } // ============ Layer 1: Sticky session ============ + // A healthy sticky account whose bounded wait queue is full may be used as a + // one-request capacity spillover in Layer 2. Keep that spillover temporary: + // rewriting the durable binding here would make a short burst migrate the + // whole conversation to a cache-cold account. + stickySpillover := false if sessionHash != "" { accountID := stickyAccountID if accountID > 0 && !isExcluded(accountID) { @@ -1215,6 +1220,7 @@ func (s *OpenAIGatewayService) selectAccountWithLoadAwareness(ctx context.Contex MaxWaiting: cfg.StickySessionMaxWaiting, }) } + stickySpillover = true } } } @@ -1366,7 +1372,7 @@ func (s *OpenAIGatewayService) selectAccountWithLoadAwareness(ctx context.Contex if selectErr != nil { return nil, true, selectErr } - if sessionHash != "" && !gatewayProfitControlGateActive(ctx) { + if sessionHash != "" && !stickySpillover && !gatewayProfitControlGateActive(ctx) { _ = s.setStickySessionAccountID(ctx, groupID, sessionHash, fresh.ID, openaiStickySessionTTL) } return selection, true, nil @@ -1405,7 +1411,7 @@ func (s *OpenAIGatewayService) selectAccountWithLoadAwareness(ctx context.Contex if selectErr != nil { return nil, selectErr } - if sessionHash != "" && !gatewayProfitControlGateActive(ctx) { + if sessionHash != "" && !stickySpillover && !gatewayProfitControlGateActive(ctx) { _ = s.setStickySessionAccountID(ctx, groupID, sessionHash, fresh.ID, openaiStickySessionTTL) } return selection, nil diff --git a/backend/internal/service/openai_gateway_service_test.go b/backend/internal/service/openai_gateway_service_test.go index 77f4ee33d6..568c948627 100644 --- a/backend/internal/service/openai_gateway_service_test.go +++ b/backend/internal/service/openai_gateway_service_test.go @@ -1207,6 +1207,49 @@ func TestOpenAISelectAccountWithLoadAwareness_StickyWaitPlan(t *testing.T) { } } +func TestOpenAISelectAccountWithLoadAwareness_StickyCapacitySpilloverKeepsBinding(t *testing.T) { + sessionHash := "sticky-spillover" + groupID := int64(1) + repo := stubOpenAIAccountRepo{ + accounts: []Account{ + {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 6, Priority: 1, GroupIDs: []int64{groupID}}, + {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 6, Priority: 1, GroupIDs: []int64{groupID}}, + }, + } + cache := &stubGatewayCache{ + sessionBindings: map[string]int64{"openai:" + sessionHash: 1}, + } + concurrencyCache := stubConcurrencyCache{ + acquireResults: map[int64]bool{1: false, 2: true}, + waitCounts: map[int64]int{1: 1}, + loadMap: map[int64]*AccountLoadInfo{ + 1: {AccountID: 1, LoadRate: 100}, + 2: {AccountID: 2, LoadRate: 10}, + }, + } + cfg := &config.Config{RunMode: config.RunModeStandard} + cfg.Gateway.Scheduling.LoadBatchEnabled = true + cfg.Gateway.Scheduling.StickySessionMaxWaiting = 1 + + svc := &OpenAIGatewayService{ + accountRepo: repo, + cache: cache, + cfg: cfg, + concurrencyService: NewConcurrencyService(concurrencyCache), + } + + selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, sessionHash, "gpt-4", nil) + require.NoError(t, err) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + require.Equal(t, int64(2), selection.Account.ID, "capacity spillover should use the other account for this request") + require.True(t, selection.Acquired) + require.Equal(t, int64(1), cache.sessionBindings["openai:"+sessionHash], "capacity spillover must not migrate the durable sticky binding") + if selection.ReleaseFunc != nil { + selection.ReleaseFunc() + } +} + func TestOpenAISelectAccountWithLoadAwareness_PrefersLowerLoad(t *testing.T) { groupID := int64(1) repo := stubOpenAIAccountRepo{