diff --git a/backend/internal/config/config.go b/backend/internal/config/config.go index 748b751463..c25d13a8b4 100644 --- a/backend/internal/config/config.go +++ b/backend/internal/config/config.go @@ -964,6 +964,8 @@ type GatewayOpenAIWSSchedulerScoreWeights struct { // Reset 倾向「会话窗口最早重置」的账号（use-it-or-lose-it）。 // >0 时，剩余重置时间越短的账号得分越高，从而被优先用尽。默认 0（关闭，不改变原有行为）。 Reset float64 `mapstructure:"reset"` + // QuotaHeadroom 倾向 7d 剩余额度更健康的账号；默认 0（关闭，不改变原有行为）。 + QuotaHeadroom float64 `mapstructure:"quota_headroom"` } // GatewayOpenAISchedulerConfig OpenAI 高级调度器配置。 @@ -1884,6 +1886,7 @@ func setDefaults() { viper.SetDefault("gateway.openai_ws.scheduler_score_weights.error_rate", 0.8) viper.SetDefault("gateway.openai_ws.scheduler_score_weights.ttft", 0.5) viper.SetDefault("gateway.openai_ws.scheduler_score_weights.reset", 0.0) + viper.SetDefault("gateway.openai_ws.scheduler_score_weights.quota_headroom", 0.0) // OpenAI HTTP upstream protocol strategy viper.SetDefault("gateway.openai_http2.enabled", true) viper.SetDefault("gateway.openai_http2.allow_proxy_fallback_to_http1", true) @@ -2663,14 +2666,16 @@ func (c *Config) Validate() error { c.Gateway.OpenAIWS.SchedulerScoreWeights.Load < 0 || c.Gateway.OpenAIWS.SchedulerScoreWeights.Queue < 0 || c.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate < 0 || - c.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT < 0 { + c.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT < 0 || + c.Gateway.OpenAIWS.SchedulerScoreWeights.QuotaHeadroom < 0 { return fmt.Errorf("gateway.openai_ws.scheduler_score_weights.* must be non-negative") } weightSum := c.Gateway.OpenAIWS.SchedulerScoreWeights.Priority + c.Gateway.OpenAIWS.SchedulerScoreWeights.Load + c.Gateway.OpenAIWS.SchedulerScoreWeights.Queue + c.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate + - c.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT + c.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT + + c.Gateway.OpenAIWS.SchedulerScoreWeights.QuotaHeadroom if weightSum <= 0 { return fmt.Errorf("gateway.openai_ws.scheduler_score_weights must not all be zero") } diff --git a/backend/internal/config/config_test.go b/backend/internal/config/config_test.go index ec0df0e3bd..be31976156 100644 --- a/backend/internal/config/config_test.go +++ b/backend/internal/config/config_test.go @@ -167,6 +167,9 @@ func TestLoadDefaultOpenAIWSConfig(t *testing.T) { if cfg.Gateway.OpenAIWS.PayloadLogSampleRate != 0.2 { t.Fatalf("Gateway.OpenAIWS.PayloadLogSampleRate = %v, want 0.2", cfg.Gateway.OpenAIWS.PayloadLogSampleRate) } + if cfg.Gateway.OpenAIWS.SchedulerScoreWeights.QuotaHeadroom != 0 { + t.Fatalf("Gateway.OpenAIWS.SchedulerScoreWeights.QuotaHeadroom = %v, want 0", cfg.Gateway.OpenAIWS.SchedulerScoreWeights.QuotaHeadroom) + } if !cfg.Gateway.OpenAIWS.StoreDisabledForceNewConn { t.Fatalf("Gateway.OpenAIWS.StoreDisabledForceNewConn = false, want true") } @@ -1717,6 +1720,11 @@ func TestValidateConfig_OpenAIWSRules(t *testing.T) { mutate: func(c *Config) { c.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = -0.1 }, wantErr: "gateway.openai_ws.scheduler_score_weights.* must be non-negative", }, + { + name: "scheduler_score_weights quota_headroom 不能为负数", + mutate: func(c *Config) { c.Gateway.OpenAIWS.SchedulerScoreWeights.QuotaHeadroom = -0.1 }, + wantErr: "gateway.openai_ws.scheduler_score_weights.* must be non-negative", + }, { name: "scheduler_score_weights 不能全为 0", mutate: func(c *Config) { @@ -1756,6 +1764,18 @@ func TestValidateConfig_OpenAIWSRules(t *testing.T) { require.Contains(t, err.Error(), tc.wantErr) }) } + + t.Run("quota_headroom 可作为唯一有效调度权重", func(t *testing.T) { + cfg := buildValid(t) + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 0 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 0 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 0 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate = 0 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT = 0 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.QuotaHeadroom = 0.1 + + require.NoError(t, cfg.Validate()) + }) } func TestValidateConfig_AutoScaleDisabledIgnoreAutoScaleFields(t *testing.T) { diff --git a/backend/internal/service/openai_account_scheduler.go b/backend/internal/service/openai_account_scheduler.go index 2277b1d904..eba93534ae 100644 --- a/backend/internal/service/openai_account_scheduler.go +++ b/backend/internal/service/openai_account_scheduler.go @@ -29,6 +29,12 @@ const ( openAIAdvancedSchedulerSettingDBTimeout = 2 * time.Second ) +const ( + openAIQuotaHeadroomNeutralFactor = 0.5 + openAIQuotaHeadroomSecondaryLowRemain = 0.10 + openAIQuotaHeadroomSnapshotStaleAfter = 8 * time.Hour +) + type cachedOpenAIAdvancedSchedulerSetting struct { enabled bool expiresAt int64 @@ -799,13 +805,18 @@ func (s *defaultOpenAIAccountScheduler) buildOpenAIAccountLoadPlan( } } } + quotaHeadroomFactor := 0.0 + if weights.QuotaHeadroom > 0 { + quotaHeadroomFactor = openAIQuotaHeadroomFactor(item.account, now) + } item.score = weights.Priority*priorityFactor + weights.Load*loadFactor + weights.Queue*queueFactor + weights.ErrorRate*errorFactor + weights.TTFT*ttftFactor + - weights.Reset*resetFactor + weights.Reset*resetFactor + + weights.QuotaHeadroom*quotaHeadroomFactor } plan.candidates = candidates @@ -1460,21 +1471,23 @@ func (s *OpenAIGatewayService) openAIStickyEscapeConfig() openAIStickyEscapeConf func (s *OpenAIGatewayService) openAIWSSchedulerWeights() GatewayOpenAIWSSchedulerScoreWeightsView { if s != nil && s.cfg != nil { return GatewayOpenAIWSSchedulerScoreWeightsView{ - Priority: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority, - Load: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load, - Queue: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue, - ErrorRate: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate, - TTFT: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT, - Reset: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Reset, + Priority: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority, + Load: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load, + Queue: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue, + ErrorRate: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate, + TTFT: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT, + Reset: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Reset, + QuotaHeadroom: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.QuotaHeadroom, } } return GatewayOpenAIWSSchedulerScoreWeightsView{ - Priority: 1.0, - Load: 1.0, - Queue: 0.7, - ErrorRate: 0.8, - TTFT: 0.5, - Reset: 0.0, + Priority: 1.0, + Load: 1.0, + Queue: 0.7, + ErrorRate: 0.8, + TTFT: 0.5, + Reset: 0.0, + QuotaHeadroom: 0.0, } } @@ -1485,7 +1498,49 @@ type GatewayOpenAIWSSchedulerScoreWeightsView struct { ErrorRate float64 TTFT float64 // Reset 倾向「会话窗口最早重置」的账号；0 表示关闭（默认）。 - Reset float64 + Reset float64 + QuotaHeadroom float64 +} + +func openAIQuotaHeadroomFactor(account *Account, now time.Time) float64 { + if account == nil || len(account.Extra) == 0 || openAIQuotaHeadroomSnapshotStale(account.Extra, now) { + return openAIQuotaHeadroomNeutralFactor + } + primaryUsedPercent, ok := resolveAccountExtraNumber(account.Extra, "codex_primary_used_percent", "codex_7d_used_percent") + if !ok || openAIQuotaWindowResetAny(account.Extra, now, "primary", "7d") { + return openAIQuotaHeadroomNeutralFactor + } + + factor := 1 - clamp01(primaryUsedPercent/100) + if secondaryUsedPercent, ok := resolveAccountExtraNumber(account.Extra, "codex_secondary_used_percent", "codex_5h_used_percent"); ok && + !openAIQuotaWindowResetAny(account.Extra, now, "secondary", "5h") { + secondaryRemaining := 1 - clamp01(secondaryUsedPercent/100) + if secondaryRemaining < openAIQuotaHeadroomSecondaryLowRemain { + factor *= openAIQuotaHeadroomNeutralFactor + } + } + return factor +} + +func openAIQuotaHeadroomSnapshotStale(extra map[string]any, now time.Time) bool { + updatedRaw, ok := extra["codex_usage_updated_at"] + if !ok { + return true + } + updatedAt, err := parseTime(fmt.Sprint(updatedRaw)) + if err != nil { + return true + } + return now.Sub(updatedAt) >= openAIQuotaHeadroomSnapshotStaleAfter +} + +func openAIQuotaWindowResetAny(extra map[string]any, now time.Time, windows ...string) bool { + for _, window := range windows { + if openAIQuotaWindowReset(extra, window, now) { + return true + } + } + return false } func clamp01(value float64) float64 { diff --git a/backend/internal/service/openai_account_scheduler_reset_test.go b/backend/internal/service/openai_account_scheduler_reset_test.go index 1ba6cb570c..50155f20d6 100644 --- a/backend/internal/service/openai_account_scheduler_reset_test.go +++ b/backend/internal/service/openai_account_scheduler_reset_test.go @@ -11,12 +11,21 @@ import ( func openAIResetTestScheduler(reset float64) *defaultOpenAIAccountScheduler { cfg := &config.Config{} cfg.Gateway.OpenAIWS.SchedulerScoreWeights = config.GatewayOpenAIWSSchedulerScoreWeights{ - Priority: 1.0, - Load: 1.0, - Queue: 0.7, - ErrorRate: 0.8, - TTFT: 0.5, - Reset: reset, + Priority: 1.0, + Load: 1.0, + Queue: 0.7, + ErrorRate: 0.8, + TTFT: 0.5, + Reset: reset, + QuotaHeadroom: 0, + } + return &defaultOpenAIAccountScheduler{service: &OpenAIGatewayService{cfg: cfg}} +} + +func openAIQuotaHeadroomTestScheduler(quotaHeadroom float64) *defaultOpenAIAccountScheduler { + cfg := &config.Config{} + cfg.Gateway.OpenAIWS.SchedulerScoreWeights = config.GatewayOpenAIWSSchedulerScoreWeights{ + QuotaHeadroom: quotaHeadroom, } return &defaultOpenAIAccountScheduler{service: &OpenAIGatewayService{cfg: cfg}} } @@ -75,3 +84,113 @@ func TestBuildOpenAIAccountLoadPlan_ResetWeightIgnoresNilWindow(t *testing.T) { scores := openAIPlanScores(plan) require.Greater(t, scores[2], scores[1], "拥有活跃窗口的账号得分高于无窗口账号") } + +func TestOpenAIQuotaHeadroomFactor_PrimaryUsedPercent(t *testing.T) { + now := time.Date(2026, 3, 11, 10, 0, 0, 0, time.UTC) + account := &Account{ + Extra: map[string]any{ + "codex_primary_used_percent": 20.0, + "codex_primary_reset_at": now.Add(24 * time.Hour).Format(time.RFC3339), + "codex_usage_updated_at": now.Add(-time.Minute).Format(time.RFC3339), + }, + } + + require.InDelta(t, 0.8, openAIQuotaHeadroomFactor(account, now), 0.0001) +} + +func TestOpenAIQuotaHeadroomFactor_PrimaryMissingIsNeutral(t *testing.T) { + now := time.Date(2026, 3, 11, 10, 0, 0, 0, time.UTC) + account := &Account{ + Extra: map[string]any{ + "codex_usage_updated_at": now.Add(-time.Minute).Format(time.RFC3339), + }, + } + + require.Equal(t, openAIQuotaHeadroomNeutralFactor, openAIQuotaHeadroomFactor(account, now)) +} + +func TestOpenAIQuotaHeadroomFactor_PrimaryResetExpiredIsNeutral(t *testing.T) { + now := time.Date(2026, 3, 11, 10, 0, 0, 0, time.UTC) + account := &Account{ + Extra: map[string]any{ + "codex_primary_used_percent": 20.0, + "codex_primary_reset_at": now.Add(-time.Minute).Format(time.RFC3339), + "codex_usage_updated_at": now.Add(-time.Minute).Format(time.RFC3339), + }, + } + + require.Equal(t, openAIQuotaHeadroomNeutralFactor, openAIQuotaHeadroomFactor(account, now)) +} + +func TestOpenAIQuotaHeadroomFactor_SecondaryLowHeadroomDiscountsPrimary(t *testing.T) { + now := time.Date(2026, 3, 11, 10, 0, 0, 0, time.UTC) + account := &Account{ + Extra: map[string]any{ + "codex_primary_used_percent": 20.0, + "codex_primary_reset_at": now.Add(24 * time.Hour).Format(time.RFC3339), + "codex_secondary_used_percent": 95.0, + "codex_secondary_reset_at": now.Add(time.Hour).Format(time.RFC3339), + "codex_usage_updated_at": now.Add(-time.Minute).Format(time.RFC3339), + }, + } + + require.InDelta(t, 0.4, openAIQuotaHeadroomFactor(account, now), 0.0001) +} + +func TestBuildOpenAIAccountLoadPlan_QuotaHeadroomPrefersHigher7dRemaining(t *testing.T) { + now := time.Now() + filtered := []*Account{ + { + ID: 1, + Priority: 0, + Extra: map[string]any{ + "codex_primary_used_percent": 80.0, + "codex_primary_reset_at": now.Add(24 * time.Hour).Format(time.RFC3339), + "codex_usage_updated_at": now.Add(-time.Minute).Format(time.RFC3339), + }, + }, + { + ID: 2, + Priority: 0, + Extra: map[string]any{ + "codex_primary_used_percent": 20.0, + "codex_primary_reset_at": now.Add(24 * time.Hour).Format(time.RFC3339), + "codex_usage_updated_at": now.Add(-time.Minute).Format(time.RFC3339), + }, + }, + } + sched := openAIQuotaHeadroomTestScheduler(1.0) + + plan := sched.buildOpenAIAccountLoadPlan(OpenAIAccountScheduleRequest{}, filtered, map[int64]*AccountLoadInfo{}) + scores := openAIPlanScores(plan) + require.Greater(t, scores[2], scores[1], "7d 剩余额度更高的账号得分应更高") +} + +func TestBuildOpenAIAccountLoadPlan_QuotaHeadroomZeroNoEffect(t *testing.T) { + now := time.Now() + filtered := []*Account{ + { + ID: 1, + Priority: 0, + Extra: map[string]any{ + "codex_primary_used_percent": 80.0, + "codex_primary_reset_at": now.Add(24 * time.Hour).Format(time.RFC3339), + "codex_usage_updated_at": now.Add(-time.Minute).Format(time.RFC3339), + }, + }, + { + ID: 2, + Priority: 0, + Extra: map[string]any{ + "codex_primary_used_percent": 20.0, + "codex_primary_reset_at": now.Add(24 * time.Hour).Format(time.RFC3339), + "codex_usage_updated_at": now.Add(-time.Minute).Format(time.RFC3339), + }, + }, + } + sched := openAIResetTestScheduler(0) + + plan := sched.buildOpenAIAccountLoadPlan(OpenAIAccountScheduleRequest{}, filtered, map[int64]*AccountLoadInfo{}) + scores := openAIPlanScores(plan) + require.Equal(t, scores[1], scores[2], "quota_headroom 权重为 0 时不应影响打分") +} diff --git a/deploy/config.example.yaml b/deploy/config.example.yaml index 67cff357e7..566c277196 100644 --- a/deploy/config.example.yaml +++ b/deploy/config.example.yaml @@ -323,6 +323,8 @@ gateway: # use-it-or-lose-it：倾向「会话窗口最早重置」的账号，剩余重置时间越短得分越高。 # 0 表示关闭（默认，不改变原有行为）；调大可让即将重置的账号被优先用尽。 reset: 0.0 + # 倾向 7d 剩余额度更健康的账号；0 表示关闭（默认，不改变原有行为），小流量灰度可设为 0.3。 + quota_headroom: 0.0 # OpenAI 高级调度器补充配置 openai_scheduler: # 是否允许 session_hash sticky 在账号健康度恶化时临时逃逸；false 可一键回退旧行为