Skip to content

Commit bf3d6c0

Browse files
committed
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.
1 parent 241023f commit bf3d6c0

12 files changed

Lines changed: 634 additions & 4 deletions

File tree

backend/internal/handler/admin/setting_handler.go

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -977,6 +977,58 @@ func (h *SettingHandler) DeleteAdminAPIKey(c *gin.Context) {
977977
response.Success(c, gin.H{"message": "Admin API key deleted"})
978978
}
979979

980+
// GetOverloadCooldownSettings 获取529过载冷却配置
981+
// GET /api/v1/admin/settings/overload-cooldown
982+
func (h *SettingHandler) GetOverloadCooldownSettings(c *gin.Context) {
983+
settings, err := h.settingService.GetOverloadCooldownSettings(c.Request.Context())
984+
if err != nil {
985+
response.ErrorFrom(c, err)
986+
return
987+
}
988+
989+
response.Success(c, dto.OverloadCooldownSettings{
990+
Enabled: settings.Enabled,
991+
CooldownMinutes: settings.CooldownMinutes,
992+
})
993+
}
994+
995+
// UpdateOverloadCooldownSettingsRequest 更新529过载冷却配置请求
996+
type UpdateOverloadCooldownSettingsRequest struct {
997+
Enabled bool `json:"enabled"`
998+
CooldownMinutes int `json:"cooldown_minutes"`
999+
}
1000+
1001+
// UpdateOverloadCooldownSettings 更新529过载冷却配置
1002+
// PUT /api/v1/admin/settings/overload-cooldown
1003+
func (h *SettingHandler) UpdateOverloadCooldownSettings(c *gin.Context) {
1004+
var req UpdateOverloadCooldownSettingsRequest
1005+
if err := c.ShouldBindJSON(&req); err != nil {
1006+
response.BadRequest(c, "Invalid request: "+err.Error())
1007+
return
1008+
}
1009+
1010+
settings := &service.OverloadCooldownSettings{
1011+
Enabled: req.Enabled,
1012+
CooldownMinutes: req.CooldownMinutes,
1013+
}
1014+
1015+
if err := h.settingService.SetOverloadCooldownSettings(c.Request.Context(), settings); err != nil {
1016+
response.BadRequest(c, err.Error())
1017+
return
1018+
}
1019+
1020+
updatedSettings, err := h.settingService.GetOverloadCooldownSettings(c.Request.Context())
1021+
if err != nil {
1022+
response.ErrorFrom(c, err)
1023+
return
1024+
}
1025+
1026+
response.Success(c, dto.OverloadCooldownSettings{
1027+
Enabled: updatedSettings.Enabled,
1028+
CooldownMinutes: updatedSettings.CooldownMinutes,
1029+
})
1030+
}
1031+
9801032
// GetStreamTimeoutSettings 获取流超时处理配置
9811033
// GET /api/v1/admin/settings/stream-timeout
9821034
func (h *SettingHandler) GetStreamTimeoutSettings(c *gin.Context) {

backend/internal/handler/dto/settings.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -157,6 +157,12 @@ type ListSoraS3ProfilesResponse struct {
157157
Items []SoraS3Profile `json:"items"`
158158
}
159159

160+
// OverloadCooldownSettings 529过载冷却配置 DTO
161+
type OverloadCooldownSettings struct {
162+
Enabled bool `json:"enabled"`
163+
CooldownMinutes int `json:"cooldown_minutes"`
164+
}
165+
160166
// StreamTimeoutSettings 流超时处理配置 DTO
161167
type StreamTimeoutSettings struct {
162168
Enabled bool `json:"enabled"`

backend/internal/server/routes/admin.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -402,6 +402,9 @@ func registerSettingsRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
402402
adminSettings.GET("/admin-api-key", h.Admin.Setting.GetAdminAPIKey)
403403
adminSettings.POST("/admin-api-key/regenerate", h.Admin.Setting.RegenerateAdminAPIKey)
404404
adminSettings.DELETE("/admin-api-key", h.Admin.Setting.DeleteAdminAPIKey)
405+
// 529过载冷却配置
406+
adminSettings.GET("/overload-cooldown", h.Admin.Setting.GetOverloadCooldownSettings)
407+
adminSettings.PUT("/overload-cooldown", h.Admin.Setting.UpdateOverloadCooldownSettings)
405408
// 流超时处理配置
406409
adminSettings.GET("/stream-timeout", h.Admin.Setting.GetStreamTimeoutSettings)
407410
adminSettings.PUT("/stream-timeout", h.Admin.Setting.UpdateStreamTimeoutSettings)

backend/internal/service/domain_constants.go

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,13 @@ const (
170170
// SettingKeyOpsRuntimeLogConfig stores JSON config for runtime log settings.
171171
SettingKeyOpsRuntimeLogConfig = "ops_runtime_log_config"
172172

173+
// =========================
174+
// Overload Cooldown (529)
175+
// =========================
176+
177+
// SettingKeyOverloadCooldownSettings stores JSON config for 529 overload cooldown handling.
178+
SettingKeyOverloadCooldownSettings = "overload_cooldown_settings"
179+
173180
// =========================
174181
// Stream Timeout Handling
175182
// =========================
Lines changed: 298 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,298 @@
1+
//go:build unit
2+
3+
package service
4+
5+
import (
6+
"context"
7+
"encoding/json"
8+
"testing"
9+
"time"
10+
11+
"github.com/Wei-Shaw/sub2api/internal/config"
12+
"github.com/stretchr/testify/require"
13+
)
14+
15+
// ---------------------------------------------------------------------------
16+
// errSettingRepo: a SettingRepository that always returns errors on read
17+
// ---------------------------------------------------------------------------
18+
19+
type errSettingRepo struct {
20+
mockSettingRepo // embed the existing mock from backup_service_test.go
21+
readErr error
22+
}
23+
24+
func (r *errSettingRepo) GetValue(_ context.Context, _ string) (string, error) {
25+
return "", r.readErr
26+
}
27+
28+
func (r *errSettingRepo) Get(_ context.Context, _ string) (*Setting, error) {
29+
return nil, r.readErr
30+
}
31+
32+
// ---------------------------------------------------------------------------
33+
// overloadAccountRepoStub: records SetOverloaded calls
34+
// ---------------------------------------------------------------------------
35+
36+
type overloadAccountRepoStub struct {
37+
mockAccountRepoForGemini
38+
overloadCalls int
39+
lastOverloadID int64
40+
lastOverloadEnd time.Time
41+
}
42+
43+
func (r *overloadAccountRepoStub) SetOverloaded(_ context.Context, id int64, until time.Time) error {
44+
r.overloadCalls++
45+
r.lastOverloadID = id
46+
r.lastOverloadEnd = until
47+
return nil
48+
}
49+
50+
// ===========================================================================
51+
// SettingService: GetOverloadCooldownSettings
52+
// ===========================================================================
53+
54+
func TestGetOverloadCooldownSettings_DefaultsWhenNotSet(t *testing.T) {
55+
repo := newMockSettingRepo()
56+
svc := NewSettingService(repo, &config.Config{})
57+
58+
settings, err := svc.GetOverloadCooldownSettings(context.Background())
59+
require.NoError(t, err)
60+
require.True(t, settings.Enabled)
61+
require.Equal(t, 10, settings.CooldownMinutes)
62+
}
63+
64+
func TestGetOverloadCooldownSettings_ReadsFromDB(t *testing.T) {
65+
repo := newMockSettingRepo()
66+
data, _ := json.Marshal(OverloadCooldownSettings{Enabled: false, CooldownMinutes: 30})
67+
repo.data[SettingKeyOverloadCooldownSettings] = string(data)
68+
svc := NewSettingService(repo, &config.Config{})
69+
70+
settings, err := svc.GetOverloadCooldownSettings(context.Background())
71+
require.NoError(t, err)
72+
require.False(t, settings.Enabled)
73+
require.Equal(t, 30, settings.CooldownMinutes)
74+
}
75+
76+
func TestGetOverloadCooldownSettings_ClampsMinValue(t *testing.T) {
77+
repo := newMockSettingRepo()
78+
data, _ := json.Marshal(OverloadCooldownSettings{Enabled: true, CooldownMinutes: 0})
79+
repo.data[SettingKeyOverloadCooldownSettings] = string(data)
80+
svc := NewSettingService(repo, &config.Config{})
81+
82+
settings, err := svc.GetOverloadCooldownSettings(context.Background())
83+
require.NoError(t, err)
84+
require.Equal(t, 1, settings.CooldownMinutes)
85+
}
86+
87+
func TestGetOverloadCooldownSettings_ClampsMaxValue(t *testing.T) {
88+
repo := newMockSettingRepo()
89+
data, _ := json.Marshal(OverloadCooldownSettings{Enabled: true, CooldownMinutes: 999})
90+
repo.data[SettingKeyOverloadCooldownSettings] = string(data)
91+
svc := NewSettingService(repo, &config.Config{})
92+
93+
settings, err := svc.GetOverloadCooldownSettings(context.Background())
94+
require.NoError(t, err)
95+
require.Equal(t, 120, settings.CooldownMinutes)
96+
}
97+
98+
func TestGetOverloadCooldownSettings_InvalidJSON_ReturnsDefaults(t *testing.T) {
99+
repo := newMockSettingRepo()
100+
repo.data[SettingKeyOverloadCooldownSettings] = "not-json"
101+
svc := NewSettingService(repo, &config.Config{})
102+
103+
settings, err := svc.GetOverloadCooldownSettings(context.Background())
104+
require.NoError(t, err)
105+
require.True(t, settings.Enabled)
106+
require.Equal(t, 10, settings.CooldownMinutes)
107+
}
108+
109+
func TestGetOverloadCooldownSettings_EmptyValue_ReturnsDefaults(t *testing.T) {
110+
repo := newMockSettingRepo()
111+
repo.data[SettingKeyOverloadCooldownSettings] = ""
112+
svc := NewSettingService(repo, &config.Config{})
113+
114+
settings, err := svc.GetOverloadCooldownSettings(context.Background())
115+
require.NoError(t, err)
116+
require.True(t, settings.Enabled)
117+
require.Equal(t, 10, settings.CooldownMinutes)
118+
}
119+
120+
// ===========================================================================
121+
// SettingService: SetOverloadCooldownSettings
122+
// ===========================================================================
123+
124+
func TestSetOverloadCooldownSettings_Success(t *testing.T) {
125+
repo := newMockSettingRepo()
126+
svc := NewSettingService(repo, &config.Config{})
127+
128+
err := svc.SetOverloadCooldownSettings(context.Background(), &OverloadCooldownSettings{
129+
Enabled: false,
130+
CooldownMinutes: 25,
131+
})
132+
require.NoError(t, err)
133+
134+
// Verify round-trip
135+
settings, err := svc.GetOverloadCooldownSettings(context.Background())
136+
require.NoError(t, err)
137+
require.False(t, settings.Enabled)
138+
require.Equal(t, 25, settings.CooldownMinutes)
139+
}
140+
141+
func TestSetOverloadCooldownSettings_RejectsNil(t *testing.T) {
142+
svc := NewSettingService(newMockSettingRepo(), &config.Config{})
143+
err := svc.SetOverloadCooldownSettings(context.Background(), nil)
144+
require.Error(t, err)
145+
}
146+
147+
func TestSetOverloadCooldownSettings_EnabledRejectsOutOfRange(t *testing.T) {
148+
svc := NewSettingService(newMockSettingRepo(), &config.Config{})
149+
150+
for _, minutes := range []int{0, -1, 121, 999} {
151+
err := svc.SetOverloadCooldownSettings(context.Background(), &OverloadCooldownSettings{
152+
Enabled: true, CooldownMinutes: minutes,
153+
})
154+
require.Error(t, err, "should reject enabled=true + cooldown_minutes=%d", minutes)
155+
require.Contains(t, err.Error(), "cooldown_minutes must be between 1-120")
156+
}
157+
}
158+
159+
func TestSetOverloadCooldownSettings_DisabledNormalizesOutOfRange(t *testing.T) {
160+
repo := newMockSettingRepo()
161+
svc := NewSettingService(repo, &config.Config{})
162+
163+
// enabled=false + cooldown_minutes=0 应该保存成功,值被归一化为10
164+
err := svc.SetOverloadCooldownSettings(context.Background(), &OverloadCooldownSettings{
165+
Enabled: false, CooldownMinutes: 0,
166+
})
167+
require.NoError(t, err, "disabled with invalid minutes should NOT be rejected")
168+
169+
// 验证持久化后读回来的值
170+
settings, err := svc.GetOverloadCooldownSettings(context.Background())
171+
require.NoError(t, err)
172+
require.False(t, settings.Enabled)
173+
require.Equal(t, 10, settings.CooldownMinutes, "should be normalized to default")
174+
}
175+
176+
func TestSetOverloadCooldownSettings_AcceptsBoundaries(t *testing.T) {
177+
svc := NewSettingService(newMockSettingRepo(), &config.Config{})
178+
179+
for _, minutes := range []int{1, 60, 120} {
180+
err := svc.SetOverloadCooldownSettings(context.Background(), &OverloadCooldownSettings{
181+
Enabled: true, CooldownMinutes: minutes,
182+
})
183+
require.NoError(t, err, "should accept cooldown_minutes=%d", minutes)
184+
}
185+
}
186+
187+
// ===========================================================================
188+
// RateLimitService: handle529 behaviour
189+
// ===========================================================================
190+
191+
func TestHandle529_EnabledFromDB_PausesAccount(t *testing.T) {
192+
accountRepo := &overloadAccountRepoStub{}
193+
settingRepo := newMockSettingRepo()
194+
data, _ := json.Marshal(OverloadCooldownSettings{Enabled: true, CooldownMinutes: 15})
195+
settingRepo.data[SettingKeyOverloadCooldownSettings] = string(data)
196+
197+
settingSvc := NewSettingService(settingRepo, &config.Config{})
198+
svc := NewRateLimitService(accountRepo, nil, &config.Config{}, nil, nil)
199+
svc.SetSettingService(settingSvc)
200+
201+
account := &Account{ID: 42, Platform: PlatformAnthropic, Type: AccountTypeOAuth}
202+
before := time.Now()
203+
svc.handle529(context.Background(), account)
204+
205+
require.Equal(t, 1, accountRepo.overloadCalls)
206+
require.Equal(t, int64(42), accountRepo.lastOverloadID)
207+
require.WithinDuration(t, before.Add(15*time.Minute), accountRepo.lastOverloadEnd, 2*time.Second)
208+
}
209+
210+
func TestHandle529_DisabledFromDB_SkipsAccount(t *testing.T) {
211+
accountRepo := &overloadAccountRepoStub{}
212+
settingRepo := newMockSettingRepo()
213+
data, _ := json.Marshal(OverloadCooldownSettings{Enabled: false, CooldownMinutes: 15})
214+
settingRepo.data[SettingKeyOverloadCooldownSettings] = string(data)
215+
216+
settingSvc := NewSettingService(settingRepo, &config.Config{})
217+
svc := NewRateLimitService(accountRepo, nil, &config.Config{}, nil, nil)
218+
svc.SetSettingService(settingSvc)
219+
220+
account := &Account{ID: 42, Platform: PlatformAnthropic, Type: AccountTypeOAuth}
221+
svc.handle529(context.Background(), account)
222+
223+
require.Equal(t, 0, accountRepo.overloadCalls, "should NOT pause when disabled")
224+
}
225+
226+
func TestHandle529_NilSettingService_FallsBackToConfig(t *testing.T) {
227+
accountRepo := &overloadAccountRepoStub{}
228+
cfg := &config.Config{}
229+
cfg.RateLimit.OverloadCooldownMinutes = 20
230+
svc := NewRateLimitService(accountRepo, nil, cfg, nil, nil)
231+
// NOT calling SetSettingService — remains nil
232+
233+
account := &Account{ID: 77, Platform: PlatformAnthropic, Type: AccountTypeOAuth}
234+
before := time.Now()
235+
svc.handle529(context.Background(), account)
236+
237+
require.Equal(t, 1, accountRepo.overloadCalls)
238+
require.WithinDuration(t, before.Add(20*time.Minute), accountRepo.lastOverloadEnd, 2*time.Second)
239+
}
240+
241+
func TestHandle529_NilSettingService_ZeroConfig_DefaultsTen(t *testing.T) {
242+
accountRepo := &overloadAccountRepoStub{}
243+
svc := NewRateLimitService(accountRepo, nil, &config.Config{}, nil, nil)
244+
245+
account := &Account{ID: 88, Platform: PlatformAnthropic, Type: AccountTypeOAuth}
246+
before := time.Now()
247+
svc.handle529(context.Background(), account)
248+
249+
require.Equal(t, 1, accountRepo.overloadCalls)
250+
require.WithinDuration(t, before.Add(10*time.Minute), accountRepo.lastOverloadEnd, 2*time.Second)
251+
}
252+
253+
func TestHandle529_DBReadError_FallsBackToConfig(t *testing.T) {
254+
accountRepo := &overloadAccountRepoStub{}
255+
errRepo := &errSettingRepo{readErr: context.DeadlineExceeded}
256+
errRepo.data = make(map[string]string)
257+
258+
cfg := &config.Config{}
259+
cfg.RateLimit.OverloadCooldownMinutes = 7
260+
settingSvc := NewSettingService(errRepo, cfg)
261+
svc := NewRateLimitService(accountRepo, nil, cfg, nil, nil)
262+
svc.SetSettingService(settingSvc)
263+
264+
account := &Account{ID: 99, Platform: PlatformAnthropic, Type: AccountTypeOAuth}
265+
before := time.Now()
266+
svc.handle529(context.Background(), account)
267+
268+
require.Equal(t, 1, accountRepo.overloadCalls)
269+
require.WithinDuration(t, before.Add(7*time.Minute), accountRepo.lastOverloadEnd, 2*time.Second)
270+
}
271+
272+
// ===========================================================================
273+
// Model: defaults & JSON round-trip
274+
// ===========================================================================
275+
276+
func TestDefaultOverloadCooldownSettings(t *testing.T) {
277+
d := DefaultOverloadCooldownSettings()
278+
require.True(t, d.Enabled)
279+
require.Equal(t, 10, d.CooldownMinutes)
280+
}
281+
282+
func TestOverloadCooldownSettings_JSONRoundTrip(t *testing.T) {
283+
original := OverloadCooldownSettings{Enabled: false, CooldownMinutes: 42}
284+
data, err := json.Marshal(original)
285+
require.NoError(t, err)
286+
287+
var decoded OverloadCooldownSettings
288+
require.NoError(t, json.Unmarshal(data, &decoded))
289+
require.Equal(t, original, decoded)
290+
291+
// Verify JSON uses snake_case field names
292+
var raw map[string]any
293+
require.NoError(t, json.Unmarshal(data, &raw))
294+
_, hasEnabled := raw["enabled"]
295+
_, hasCooldown := raw["cooldown_minutes"]
296+
require.True(t, hasEnabled, "JSON must use 'enabled'")
297+
require.True(t, hasCooldown, "JSON must use 'cooldown_minutes'")
298+
}

0 commit comments

Comments
 (0)