feat: add 529 overload cooldown toggle and duration settings in admin gateway page

Move 529 overload cooldown configuration from config file to admin
settings UI. Adds an enable/disable toggle and configurable cooldown
duration (1-120 min) under /admin/settings gateway tab, stored as
JSON in the settings table.

When disabled, 529 errors are logged but accounts are no longer
paused from scheduling. Falls back to config file value when DB
is unreachable or settingService is nil.
This commit is contained in:
shaw
2026-03-18 16:22:19 +08:00
parent 241023f3fc
commit bf3d6c0e6e
12 changed files with 634 additions and 4 deletions

View File

@@ -977,6 +977,58 @@ func (h *SettingHandler) DeleteAdminAPIKey(c *gin.Context) {
response.Success(c, gin.H{"message": "Admin API key deleted"})
}
// GetOverloadCooldownSettings 获取529过载冷却配置
// GET /api/v1/admin/settings/overload-cooldown
func (h *SettingHandler) GetOverloadCooldownSettings(c *gin.Context) {
settings, err := h.settingService.GetOverloadCooldownSettings(c.Request.Context())
if err != nil {
response.ErrorFrom(c, err)
return
}
response.Success(c, dto.OverloadCooldownSettings{
Enabled: settings.Enabled,
CooldownMinutes: settings.CooldownMinutes,
})
}
// UpdateOverloadCooldownSettingsRequest 更新529过载冷却配置请求
type UpdateOverloadCooldownSettingsRequest struct {
Enabled bool `json:"enabled"`
CooldownMinutes int `json:"cooldown_minutes"`
}
// UpdateOverloadCooldownSettings 更新529过载冷却配置
// PUT /api/v1/admin/settings/overload-cooldown
func (h *SettingHandler) UpdateOverloadCooldownSettings(c *gin.Context) {
var req UpdateOverloadCooldownSettingsRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.BadRequest(c, "Invalid request: "+err.Error())
return
}
settings := &service.OverloadCooldownSettings{
Enabled: req.Enabled,
CooldownMinutes: req.CooldownMinutes,
}
if err := h.settingService.SetOverloadCooldownSettings(c.Request.Context(), settings); err != nil {
response.BadRequest(c, err.Error())
return
}
updatedSettings, err := h.settingService.GetOverloadCooldownSettings(c.Request.Context())
if err != nil {
response.ErrorFrom(c, err)
return
}
response.Success(c, dto.OverloadCooldownSettings{
Enabled: updatedSettings.Enabled,
CooldownMinutes: updatedSettings.CooldownMinutes,
})
}
// GetStreamTimeoutSettings 获取流超时处理配置
// GET /api/v1/admin/settings/stream-timeout
func (h *SettingHandler) GetStreamTimeoutSettings(c *gin.Context) {

View File

@@ -157,6 +157,12 @@ type ListSoraS3ProfilesResponse struct {
Items []SoraS3Profile `json:"items"`
}
// OverloadCooldownSettings 529过载冷却配置 DTO
type OverloadCooldownSettings struct {
Enabled bool `json:"enabled"`
CooldownMinutes int `json:"cooldown_minutes"`
}
// StreamTimeoutSettings 流超时处理配置 DTO
type StreamTimeoutSettings struct {
Enabled bool `json:"enabled"`

View File

@@ -402,6 +402,9 @@ func registerSettingsRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
adminSettings.GET("/admin-api-key", h.Admin.Setting.GetAdminAPIKey)
adminSettings.POST("/admin-api-key/regenerate", h.Admin.Setting.RegenerateAdminAPIKey)
adminSettings.DELETE("/admin-api-key", h.Admin.Setting.DeleteAdminAPIKey)
// 529过载冷却配置
adminSettings.GET("/overload-cooldown", h.Admin.Setting.GetOverloadCooldownSettings)
adminSettings.PUT("/overload-cooldown", h.Admin.Setting.UpdateOverloadCooldownSettings)
// 流超时处理配置
adminSettings.GET("/stream-timeout", h.Admin.Setting.GetStreamTimeoutSettings)
adminSettings.PUT("/stream-timeout", h.Admin.Setting.UpdateStreamTimeoutSettings)

View File

@@ -170,6 +170,13 @@ const (
// SettingKeyOpsRuntimeLogConfig stores JSON config for runtime log settings.
SettingKeyOpsRuntimeLogConfig = "ops_runtime_log_config"
// =========================
// Overload Cooldown (529)
// =========================
// SettingKeyOverloadCooldownSettings stores JSON config for 529 overload cooldown handling.
SettingKeyOverloadCooldownSettings = "overload_cooldown_settings"
// =========================
// Stream Timeout Handling
// =========================

View File

@@ -0,0 +1,298 @@
//go:build unit
package service
import (
"context"
"encoding/json"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/stretchr/testify/require"
)
// ---------------------------------------------------------------------------
// errSettingRepo: a SettingRepository that always returns errors on read
// ---------------------------------------------------------------------------
type errSettingRepo struct {
mockSettingRepo // embed the existing mock from backup_service_test.go
readErr error
}
func (r *errSettingRepo) GetValue(_ context.Context, _ string) (string, error) {
return "", r.readErr
}
func (r *errSettingRepo) Get(_ context.Context, _ string) (*Setting, error) {
return nil, r.readErr
}
// ---------------------------------------------------------------------------
// overloadAccountRepoStub: records SetOverloaded calls
// ---------------------------------------------------------------------------
type overloadAccountRepoStub struct {
mockAccountRepoForGemini
overloadCalls int
lastOverloadID int64
lastOverloadEnd time.Time
}
func (r *overloadAccountRepoStub) SetOverloaded(_ context.Context, id int64, until time.Time) error {
r.overloadCalls++
r.lastOverloadID = id
r.lastOverloadEnd = until
return nil
}
// ===========================================================================
// SettingService: GetOverloadCooldownSettings
// ===========================================================================
func TestGetOverloadCooldownSettings_DefaultsWhenNotSet(t *testing.T) {
repo := newMockSettingRepo()
svc := NewSettingService(repo, &config.Config{})
settings, err := svc.GetOverloadCooldownSettings(context.Background())
require.NoError(t, err)
require.True(t, settings.Enabled)
require.Equal(t, 10, settings.CooldownMinutes)
}
func TestGetOverloadCooldownSettings_ReadsFromDB(t *testing.T) {
repo := newMockSettingRepo()
data, _ := json.Marshal(OverloadCooldownSettings{Enabled: false, CooldownMinutes: 30})
repo.data[SettingKeyOverloadCooldownSettings] = string(data)
svc := NewSettingService(repo, &config.Config{})
settings, err := svc.GetOverloadCooldownSettings(context.Background())
require.NoError(t, err)
require.False(t, settings.Enabled)
require.Equal(t, 30, settings.CooldownMinutes)
}
func TestGetOverloadCooldownSettings_ClampsMinValue(t *testing.T) {
repo := newMockSettingRepo()
data, _ := json.Marshal(OverloadCooldownSettings{Enabled: true, CooldownMinutes: 0})
repo.data[SettingKeyOverloadCooldownSettings] = string(data)
svc := NewSettingService(repo, &config.Config{})
settings, err := svc.GetOverloadCooldownSettings(context.Background())
require.NoError(t, err)
require.Equal(t, 1, settings.CooldownMinutes)
}
func TestGetOverloadCooldownSettings_ClampsMaxValue(t *testing.T) {
repo := newMockSettingRepo()
data, _ := json.Marshal(OverloadCooldownSettings{Enabled: true, CooldownMinutes: 999})
repo.data[SettingKeyOverloadCooldownSettings] = string(data)
svc := NewSettingService(repo, &config.Config{})
settings, err := svc.GetOverloadCooldownSettings(context.Background())
require.NoError(t, err)
require.Equal(t, 120, settings.CooldownMinutes)
}
func TestGetOverloadCooldownSettings_InvalidJSON_ReturnsDefaults(t *testing.T) {
repo := newMockSettingRepo()
repo.data[SettingKeyOverloadCooldownSettings] = "not-json"
svc := NewSettingService(repo, &config.Config{})
settings, err := svc.GetOverloadCooldownSettings(context.Background())
require.NoError(t, err)
require.True(t, settings.Enabled)
require.Equal(t, 10, settings.CooldownMinutes)
}
func TestGetOverloadCooldownSettings_EmptyValue_ReturnsDefaults(t *testing.T) {
repo := newMockSettingRepo()
repo.data[SettingKeyOverloadCooldownSettings] = ""
svc := NewSettingService(repo, &config.Config{})
settings, err := svc.GetOverloadCooldownSettings(context.Background())
require.NoError(t, err)
require.True(t, settings.Enabled)
require.Equal(t, 10, settings.CooldownMinutes)
}
// ===========================================================================
// SettingService: SetOverloadCooldownSettings
// ===========================================================================
func TestSetOverloadCooldownSettings_Success(t *testing.T) {
repo := newMockSettingRepo()
svc := NewSettingService(repo, &config.Config{})
err := svc.SetOverloadCooldownSettings(context.Background(), &OverloadCooldownSettings{
Enabled: false,
CooldownMinutes: 25,
})
require.NoError(t, err)
// Verify round-trip
settings, err := svc.GetOverloadCooldownSettings(context.Background())
require.NoError(t, err)
require.False(t, settings.Enabled)
require.Equal(t, 25, settings.CooldownMinutes)
}
func TestSetOverloadCooldownSettings_RejectsNil(t *testing.T) {
svc := NewSettingService(newMockSettingRepo(), &config.Config{})
err := svc.SetOverloadCooldownSettings(context.Background(), nil)
require.Error(t, err)
}
func TestSetOverloadCooldownSettings_EnabledRejectsOutOfRange(t *testing.T) {
svc := NewSettingService(newMockSettingRepo(), &config.Config{})
for _, minutes := range []int{0, -1, 121, 999} {
err := svc.SetOverloadCooldownSettings(context.Background(), &OverloadCooldownSettings{
Enabled: true, CooldownMinutes: minutes,
})
require.Error(t, err, "should reject enabled=true + cooldown_minutes=%d", minutes)
require.Contains(t, err.Error(), "cooldown_minutes must be between 1-120")
}
}
func TestSetOverloadCooldownSettings_DisabledNormalizesOutOfRange(t *testing.T) {
repo := newMockSettingRepo()
svc := NewSettingService(repo, &config.Config{})
// enabled=false + cooldown_minutes=0 应该保存成功值被归一化为10
err := svc.SetOverloadCooldownSettings(context.Background(), &OverloadCooldownSettings{
Enabled: false, CooldownMinutes: 0,
})
require.NoError(t, err, "disabled with invalid minutes should NOT be rejected")
// 验证持久化后读回来的值
settings, err := svc.GetOverloadCooldownSettings(context.Background())
require.NoError(t, err)
require.False(t, settings.Enabled)
require.Equal(t, 10, settings.CooldownMinutes, "should be normalized to default")
}
func TestSetOverloadCooldownSettings_AcceptsBoundaries(t *testing.T) {
svc := NewSettingService(newMockSettingRepo(), &config.Config{})
for _, minutes := range []int{1, 60, 120} {
err := svc.SetOverloadCooldownSettings(context.Background(), &OverloadCooldownSettings{
Enabled: true, CooldownMinutes: minutes,
})
require.NoError(t, err, "should accept cooldown_minutes=%d", minutes)
}
}
// ===========================================================================
// RateLimitService: handle529 behaviour
// ===========================================================================
func TestHandle529_EnabledFromDB_PausesAccount(t *testing.T) {
accountRepo := &overloadAccountRepoStub{}
settingRepo := newMockSettingRepo()
data, _ := json.Marshal(OverloadCooldownSettings{Enabled: true, CooldownMinutes: 15})
settingRepo.data[SettingKeyOverloadCooldownSettings] = string(data)
settingSvc := NewSettingService(settingRepo, &config.Config{})
svc := NewRateLimitService(accountRepo, nil, &config.Config{}, nil, nil)
svc.SetSettingService(settingSvc)
account := &Account{ID: 42, Platform: PlatformAnthropic, Type: AccountTypeOAuth}
before := time.Now()
svc.handle529(context.Background(), account)
require.Equal(t, 1, accountRepo.overloadCalls)
require.Equal(t, int64(42), accountRepo.lastOverloadID)
require.WithinDuration(t, before.Add(15*time.Minute), accountRepo.lastOverloadEnd, 2*time.Second)
}
func TestHandle529_DisabledFromDB_SkipsAccount(t *testing.T) {
accountRepo := &overloadAccountRepoStub{}
settingRepo := newMockSettingRepo()
data, _ := json.Marshal(OverloadCooldownSettings{Enabled: false, CooldownMinutes: 15})
settingRepo.data[SettingKeyOverloadCooldownSettings] = string(data)
settingSvc := NewSettingService(settingRepo, &config.Config{})
svc := NewRateLimitService(accountRepo, nil, &config.Config{}, nil, nil)
svc.SetSettingService(settingSvc)
account := &Account{ID: 42, Platform: PlatformAnthropic, Type: AccountTypeOAuth}
svc.handle529(context.Background(), account)
require.Equal(t, 0, accountRepo.overloadCalls, "should NOT pause when disabled")
}
func TestHandle529_NilSettingService_FallsBackToConfig(t *testing.T) {
accountRepo := &overloadAccountRepoStub{}
cfg := &config.Config{}
cfg.RateLimit.OverloadCooldownMinutes = 20
svc := NewRateLimitService(accountRepo, nil, cfg, nil, nil)
// NOT calling SetSettingService — remains nil
account := &Account{ID: 77, Platform: PlatformAnthropic, Type: AccountTypeOAuth}
before := time.Now()
svc.handle529(context.Background(), account)
require.Equal(t, 1, accountRepo.overloadCalls)
require.WithinDuration(t, before.Add(20*time.Minute), accountRepo.lastOverloadEnd, 2*time.Second)
}
func TestHandle529_NilSettingService_ZeroConfig_DefaultsTen(t *testing.T) {
accountRepo := &overloadAccountRepoStub{}
svc := NewRateLimitService(accountRepo, nil, &config.Config{}, nil, nil)
account := &Account{ID: 88, Platform: PlatformAnthropic, Type: AccountTypeOAuth}
before := time.Now()
svc.handle529(context.Background(), account)
require.Equal(t, 1, accountRepo.overloadCalls)
require.WithinDuration(t, before.Add(10*time.Minute), accountRepo.lastOverloadEnd, 2*time.Second)
}
func TestHandle529_DBReadError_FallsBackToConfig(t *testing.T) {
accountRepo := &overloadAccountRepoStub{}
errRepo := &errSettingRepo{readErr: context.DeadlineExceeded}
errRepo.data = make(map[string]string)
cfg := &config.Config{}
cfg.RateLimit.OverloadCooldownMinutes = 7
settingSvc := NewSettingService(errRepo, cfg)
svc := NewRateLimitService(accountRepo, nil, cfg, nil, nil)
svc.SetSettingService(settingSvc)
account := &Account{ID: 99, Platform: PlatformAnthropic, Type: AccountTypeOAuth}
before := time.Now()
svc.handle529(context.Background(), account)
require.Equal(t, 1, accountRepo.overloadCalls)
require.WithinDuration(t, before.Add(7*time.Minute), accountRepo.lastOverloadEnd, 2*time.Second)
}
// ===========================================================================
// Model: defaults & JSON round-trip
// ===========================================================================
func TestDefaultOverloadCooldownSettings(t *testing.T) {
d := DefaultOverloadCooldownSettings()
require.True(t, d.Enabled)
require.Equal(t, 10, d.CooldownMinutes)
}
func TestOverloadCooldownSettings_JSONRoundTrip(t *testing.T) {
original := OverloadCooldownSettings{Enabled: false, CooldownMinutes: 42}
data, err := json.Marshal(original)
require.NoError(t, err)
var decoded OverloadCooldownSettings
require.NoError(t, json.Unmarshal(data, &decoded))
require.Equal(t, original, decoded)
// Verify JSON uses snake_case field names
var raw map[string]any
require.NoError(t, json.Unmarshal(data, &raw))
_, hasEnabled := raw["enabled"]
_, hasCooldown := raw["cooldown_minutes"]
require.True(t, hasEnabled, "JSON must use 'enabled'")
require.True(t, hasCooldown, "JSON must use 'cooldown_minutes'")
}

View File

@@ -1023,11 +1023,34 @@ func parseOpenAIRateLimitResetTime(body []byte) *int64 {
}
// handle529 处理529过载错误
// 根据配置设置过载冷却时
// 根据配置决定是否暂停账号调度及冷却时
func (s *RateLimitService) handle529(ctx context.Context, account *Account) {
cooldownMinutes := s.cfg.RateLimit.OverloadCooldownMinutes
var settings *OverloadCooldownSettings
if s.settingService != nil {
var err error
settings, err = s.settingService.GetOverloadCooldownSettings(ctx)
if err != nil {
slog.Warn("overload_settings_read_failed", "account_id", account.ID, "error", err)
settings = nil
}
}
// 回退到配置文件
if settings == nil {
cooldown := s.cfg.RateLimit.OverloadCooldownMinutes
if cooldown <= 0 {
cooldown = 10
}
settings = &OverloadCooldownSettings{Enabled: true, CooldownMinutes: cooldown}
}
if !settings.Enabled {
slog.Info("account_529_ignored", "account_id", account.ID, "reason", "overload_cooldown_disabled")
return
}
cooldownMinutes := settings.CooldownMinutes
if cooldownMinutes <= 0 {
cooldownMinutes = 10 // 默认10分钟
cooldownMinutes = 10
}
until := time.Now().Add(time.Duration(cooldownMinutes) * time.Minute)

View File

@@ -1172,6 +1172,57 @@ func (s *SettingService) GetLinuxDoConnectOAuthConfig(ctx context.Context) (conf
return effective, nil
}
// GetOverloadCooldownSettings 获取529过载冷却配置
func (s *SettingService) GetOverloadCooldownSettings(ctx context.Context) (*OverloadCooldownSettings, error) {
value, err := s.settingRepo.GetValue(ctx, SettingKeyOverloadCooldownSettings)
if err != nil {
if errors.Is(err, ErrSettingNotFound) {
return DefaultOverloadCooldownSettings(), nil
}
return nil, fmt.Errorf("get overload cooldown settings: %w", err)
}
if value == "" {
return DefaultOverloadCooldownSettings(), nil
}
var settings OverloadCooldownSettings
if err := json.Unmarshal([]byte(value), &settings); err != nil {
return DefaultOverloadCooldownSettings(), nil
}
// 修正配置值范围
if settings.CooldownMinutes < 1 {
settings.CooldownMinutes = 1
}
if settings.CooldownMinutes > 120 {
settings.CooldownMinutes = 120
}
return &settings, nil
}
// SetOverloadCooldownSettings 设置529过载冷却配置
func (s *SettingService) SetOverloadCooldownSettings(ctx context.Context, settings *OverloadCooldownSettings) error {
if settings == nil {
return fmt.Errorf("settings cannot be nil")
}
// 禁用时修正为合法值即可,不拒绝请求
if settings.CooldownMinutes < 1 || settings.CooldownMinutes > 120 {
if settings.Enabled {
return fmt.Errorf("cooldown_minutes must be between 1-120")
}
settings.CooldownMinutes = 10 // 禁用状态下归一化为默认值
}
data, err := json.Marshal(settings)
if err != nil {
return fmt.Errorf("marshal overload cooldown settings: %w", err)
}
return s.settingRepo.Set(ctx, SettingKeyOverloadCooldownSettings, string(data))
}
// GetStreamTimeoutSettings 获取流超时处理配置
func (s *SettingService) GetStreamTimeoutSettings(ctx context.Context) (*StreamTimeoutSettings, error) {
value, err := s.settingRepo.GetValue(ctx, SettingKeyStreamTimeoutSettings)

View File

@@ -222,6 +222,22 @@ type BetaPolicySettings struct {
Rules []BetaPolicyRule `json:"rules"`
}
// OverloadCooldownSettings 529过载冷却配置
type OverloadCooldownSettings struct {
// Enabled 是否在收到529时暂停账号调度
Enabled bool `json:"enabled"`
// CooldownMinutes 冷却时长(分钟)
CooldownMinutes int `json:"cooldown_minutes"`
}
// DefaultOverloadCooldownSettings 返回默认的过载冷却配置启用10分钟
func DefaultOverloadCooldownSettings() *OverloadCooldownSettings {
return &OverloadCooldownSettings{
Enabled: true,
CooldownMinutes: 10,
}
}
// DefaultBetaPolicySettings 返回默认的 Beta 策略配置
func DefaultBetaPolicySettings() *BetaPolicySettings {
return &BetaPolicySettings{