Skip to content

Commit 9fd95df

Browse files
Wei-ShawWei-Shaw
authored
Merge pull request Wei-Shaw#679 from DaydreamCoding/feat/account-rpm-limit
feat: 添加账号级别 RPM（每分钟请求数）限流功能
2 parents 54de3bf + 212cbbd commit 9fd95df

27 files changed

Lines changed: 1174 additions & 31 deletions

‎backend/cmd/server/wire_gen.go‎

Lines changed: 3 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎backend/internal/handler/admin/account_data_handler_test.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ func setupAccountDataRouter() (*gin.Engine, *stubAdminService) {
6464
nil,
6565
nil,
6666
nil,
67+
nil,
6768
)
6869

6970
router.GET("/api/v1/admin/accounts/data", h.ExportData)

‎backend/internal/handler/admin/account_handler.go‎

Lines changed: 60 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ type AccountHandler struct {
5353
concurrencyService *service.ConcurrencyService
5454
crsSyncService *service.CRSSyncService
5555
sessionLimitCache service.SessionLimitCache
56+
rpmCache service.RPMCache
5657
tokenCacheInvalidator service.TokenCacheInvalidator
5758
}
5859

@@ -69,6 +70,7 @@ func NewAccountHandler(
6970
concurrencyService *service.ConcurrencyService,
7071
crsSyncService *service.CRSSyncService,
7172
sessionLimitCache service.SessionLimitCache,
73+
rpmCache service.RPMCache,
7274
tokenCacheInvalidator service.TokenCacheInvalidator,
7375
) *AccountHandler {
7476
return &AccountHandler{
@@ -83,6 +85,7 @@ func NewAccountHandler(
8385
concurrencyService: concurrencyService,
8486
crsSyncService: crsSyncService,
8587
sessionLimitCache: sessionLimitCache,
88+
rpmCache: rpmCache,
8689
tokenCacheInvalidator: tokenCacheInvalidator,
8790
}
8891
}
@@ -154,6 +157,7 @@ type AccountWithConcurrency struct {
154157
// 以下字段仅对 Anthropic OAuth/SetupToken 账号有效，且仅在启用相应功能时返回
155158
CurrentWindowCost *float64 `json:"current_window_cost,omitempty"` // 当前窗口费用
156159
ActiveSessions *int `json:"active_sessions,omitempty"` // 当前活跃会话数
160+
CurrentRPM *int `json:"current_rpm,omitempty"` // 当前分钟 RPM 计数
157161
}
158162

159163
func (h *AccountHandler) buildAccountResponseWithRuntime(ctx context.Context, account *service.Account) AccountWithConcurrency {
@@ -189,6 +193,12 @@ func (h *AccountHandler) buildAccountResponseWithRuntime(ctx context.Context, ac
189193
}
190194
}
191195
}
196+
197+
if h.rpmCache != nil && account.GetBaseRPM() > 0 {
198+
if rpm, err := h.rpmCache.GetRPM(ctx, account.ID); err == nil {
199+
item.CurrentRPM = &rpm
200+
}
201+
}
192202
}
193203

194204
return item
@@ -231,9 +241,10 @@ func (h *AccountHandler) List(c *gin.Context) {
231241
concurrencyCounts = make(map[int64]int)
232242
}
233243

234-
// 识别需要查询窗口费用和会话数的账号（Anthropic OAuth/SetupToken 且启用了相应功能）
244+
// 识别需要查询窗口费用、会话数和 RPM 的账号（Anthropic OAuth/SetupToken 且启用了相应功能）
235245
windowCostAccountIDs := make([]int64, 0)
236246
sessionLimitAccountIDs := make([]int64, 0)
247+
rpmAccountIDs := make([]int64, 0)
237248
sessionIdleTimeouts := make(map[int64]time.Duration) // 各账号的会话空闲超时配置
238249
for i := range accounts {
239250
acc := &accounts[i]
@@ -245,12 +256,24 @@ func (h *AccountHandler) List(c *gin.Context) {
245256
sessionLimitAccountIDs = append(sessionLimitAccountIDs, acc.ID)
246257
sessionIdleTimeouts[acc.ID] = time.Duration(acc.GetSessionIdleTimeoutMinutes()) * time.Minute
247258
}
259+
if acc.GetBaseRPM() > 0 {
260+
rpmAccountIDs = append(rpmAccountIDs, acc.ID)
261+
}
248262
}
249263
}
250264

251-
// 并行获取窗口费用和活跃会话数
265+
// 并行获取窗口费用、活跃会话数和 RPM 计数
252266
var windowCosts map[int64]float64
253267
var activeSessions map[int64]int
268+
var rpmCounts map[int64]int
269+
270+
// 获取 RPM 计数（批量查询）
271+
if len(rpmAccountIDs) > 0 && h.rpmCache != nil {
272+
rpmCounts, _ = h.rpmCache.GetRPMBatch(c.Request.Context(), rpmAccountIDs)
273+
if rpmCounts == nil {
274+
rpmCounts = make(map[int64]int)
275+
}
276+
}
254277

255278
// 获取活跃会话数（批量查询，传入各账号的 idleTimeout 配置）
256279
if len(sessionLimitAccountIDs) > 0 && h.sessionLimitCache != nil {
@@ -311,6 +334,13 @@ func (h *AccountHandler) List(c *gin.Context) {
311334
}
312335
}
313336

337+
// 添加 RPM 计数（仅当启用时）
338+
if rpmCounts != nil {
339+
if rpm, ok := rpmCounts[acc.ID]; ok {
340+
item.CurrentRPM = &rpm
341+
}
342+
}
343+
314344
result[i] = item
315345
}
316346

@@ -453,6 +483,8 @@ func (h *AccountHandler) Create(c *gin.Context) {
453483
response.BadRequest(c, "rate_multiplier must be >= 0")
454484
return
455485
}
486+
// base_rpm 输入校验：负值归零，超过 10000 截断
487+
sanitizeExtraBaseRPM(req.Extra)
456488

457489
// 确定是否跳过混合渠道检查
458490
skipCheck := req.ConfirmMixedChannelRisk != nil && *req.ConfirmMixedChannelRisk
@@ -522,6 +554,8 @@ func (h *AccountHandler) Update(c *gin.Context) {
522554
response.BadRequest(c, "rate_multiplier must be >= 0")
523555
return
524556
}
557+
// base_rpm 输入校验：负值归零，超过 10000 截断
558+
sanitizeExtraBaseRPM(req.Extra)
525559

526560
// 确定是否跳过混合渠道检查
527561
skipCheck := req.ConfirmMixedChannelRisk != nil && *req.ConfirmMixedChannelRisk
@@ -904,6 +938,9 @@ func (h *AccountHandler) BatchCreate(c *gin.Context) {
904938
continue
905939
}
906940

941+
// base_rpm 输入校验：负值归零，超过 10000 截断
942+
sanitizeExtraBaseRPM(item.Extra)
943+
907944
skipCheck := item.ConfirmMixedChannelRisk != nil && *item.ConfirmMixedChannelRisk
908945

909946
account, err := h.adminService.CreateAccount(ctx, &service.CreateAccountInput{
@@ -1048,6 +1085,8 @@ func (h *AccountHandler) BulkUpdate(c *gin.Context) {
10481085
response.BadRequest(c, "rate_multiplier must be >= 0")
10491086
return
10501087
}
1088+
// base_rpm 输入校验：负值归零，超过 10000 截断
1089+
sanitizeExtraBaseRPM(req.Extra)
10511090

10521091
// 确定是否跳过混合渠道检查
10531092
skipCheck := req.ConfirmMixedChannelRisk != nil && *req.ConfirmMixedChannelRisk
@@ -1706,3 +1745,22 @@ func (h *AccountHandler) BatchRefreshTier(c *gin.Context) {
17061745
func (h *AccountHandler) GetAntigravityDefaultModelMapping(c *gin.Context) {
17071746
response.Success(c, domain.DefaultAntigravityModelMapping)
17081747
}
1748+
1749+
// sanitizeExtraBaseRPM 对 extra map 中的 base_rpm 值进行范围校验和归一化。
1750+
// 负值归零，超过 10000 截断为 10000。extra 为 nil 或不含 base_rpm 时无操作。
1751+
func sanitizeExtraBaseRPM(extra map[string]any) {
1752+
if extra == nil {
1753+
return
1754+
}
1755+
raw, ok := extra["base_rpm"]
1756+
if !ok {
1757+
return
1758+
}
1759+
v := service.ParseExtraInt(raw)
1760+
if v < 0 {
1761+
v = 0
1762+
} else if v > 10000 {
1763+
v = 10000
1764+
}
1765+
extra["base_rpm"] = v
1766+
}

‎backend/internal/handler/admin/account_handler_mixed_channel_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ import (
1515
func setupAccountMixedChannelRouter(adminSvc *stubAdminService) *gin.Engine {
1616
gin.SetMode(gin.TestMode)
1717
router := gin.New()
18-
accountHandler := NewAccountHandler(adminSvc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
18+
accountHandler := NewAccountHandler(adminSvc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
1919
router.POST("/api/v1/admin/accounts/check-mixed-channel", accountHandler.CheckMixedChannel)
2020
router.POST("/api/v1/admin/accounts", accountHandler.Create)
2121
router.PUT("/api/v1/admin/accounts/:id", accountHandler.Update)

‎backend/internal/handler/admin/account_handler_passthrough_test.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ func TestAccountHandler_Create_AnthropicAPIKeyPassthroughExtraForwarded(t *testi
2828
nil,
2929
nil,
3030
nil,
31+
nil,
3132
)
3233

3334
router := gin.New()

‎backend/internal/handler/admin/batch_update_credentials_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ func (f *failingAdminService) UpdateAccount(ctx context.Context, id int64, input
3636
func setupAccountHandlerWithService(adminSvc service.AdminService) (*gin.Engine, *AccountHandler) {
3737
gin.SetMode(gin.TestMode)
3838
router := gin.New()
39-
handler := NewAccountHandler(adminSvc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
39+
handler := NewAccountHandler(adminSvc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
4040
router.POST("/api/v1/admin/accounts/batch-update-credentials", handler.BatchUpdateCredentials)
4141
return router, handler
4242
}

‎backend/internal/handler/dto/mappers.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,13 @@ func AccountFromServiceShallow(a *service.Account) *Account {
209209
if idleTimeout := a.GetSessionIdleTimeoutMinutes(); idleTimeout > 0 {
210210
out.SessionIdleTimeoutMin = &idleTimeout
211211
}
212+
if rpm := a.GetBaseRPM(); rpm > 0 {
213+
out.BaseRPM = &rpm
214+
strategy := a.GetRPMStrategy()
215+
out.RPMStrategy = &strategy
216+
buffer := a.GetRPMStickyBuffer()
217+
out.RPMStickyBuffer = &buffer
218+
}
212219
// TLS指纹伪装开关
213220
if a.IsTLSFingerprintEnabled() {
214221
enabled := true

‎backend/internal/handler/dto/types.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,6 +153,12 @@ type Account struct {
153153
MaxSessions *int `json:"max_sessions,omitempty"`
154154
SessionIdleTimeoutMin *int `json:"session_idle_timeout_minutes,omitempty"`
155155

156+
// RPM 限制（仅 Anthropic OAuth/SetupToken 账号有效）
157+
// 从 extra 字段提取，方便前端显示和编辑
158+
BaseRPM *int `json:"base_rpm,omitempty"`
159+
RPMStrategy *string `json:"rpm_strategy,omitempty"`
160+
RPMStickyBuffer *int `json:"rpm_sticky_buffer,omitempty"`
161+
156162
// TLS指纹伪装（仅 Anthropic OAuth/SetupToken 账号有效）
157163
// 从 extra 字段提取，方便前端显示和编辑
158164
EnableTLSFingerprint *bool `json:"enable_tls_fingerprint,omitempty"`

‎backend/internal/handler/gateway_handler.go‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -403,6 +403,15 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
403403
return
404404
}
405405

406+
// RPM 计数递增（Forward 成功后）
407+
// 注意：TOCTOU 竞态是已知且可接受的设计权衡，与 WindowCost 一致的 soft-limit 模式。
408+
// 在高并发下可能短暂超出 RPM 限制，但不会导致请求失败。
409+
if account.IsAnthropicOAuthOrSetupToken() && account.GetBaseRPM() > 0 {
410+
if err := h.gatewayService.IncrementAccountRPM(c.Request.Context(), account.ID); err != nil {
411+
reqLog.Warn("gateway.rpm_increment_failed", zap.Int64("account_id", account.ID), zap.Error(err))
412+
}
413+
}
414+
406415
// 捕获请求信息（用于异步记录，避免在 goroutine 中访问 gin.Context）
407416
userAgent := c.GetHeader("User-Agent")
408417
clientIP := ip.GetClientIP(c)
@@ -595,7 +604,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
595604
h.handleStreamingAwareError(c, status, code, message, streamStarted)
596605
return
597606
}
598-
// 兜底重试按“直接请求兜底分组”处理：清除强制平台，允许按分组平台调度
607+
// 兜底重试按"直接请求兜底分组"处理：清除强制平台，允许按分组平台调度
599608
ctx := context.WithValue(c.Request.Context(), ctxkey.ForcePlatform, "")
600609
c.Request = c.Request.WithContext(ctx)
601610
currentAPIKey = fallbackAPIKey
@@ -629,6 +638,15 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
629638
return
630639
}
631640

641+
// RPM 计数递增（Forward 成功后）
642+
// 注意：TOCTOU 竞态是已知且可接受的设计权衡，与 WindowCost 一致的 soft-limit 模式。
643+
// 在高并发下可能短暂超出 RPM 限制，但不会导致请求失败。
644+
if account.IsAnthropicOAuthOrSetupToken() && account.GetBaseRPM() > 0 {
645+
if err := h.gatewayService.IncrementAccountRPM(c.Request.Context(), account.ID); err != nil {
646+
reqLog.Warn("gateway.rpm_increment_failed", zap.Int64("account_id", account.ID), zap.Error(err))
647+
}
648+
}
649+
632650
// 捕获请求信息（用于异步记录，避免在 goroutine 中访问 gin.Context）
633651
userAgent := c.GetHeader("User-Agent")
634652
clientIP := ip.GetClientIP(c)

‎backend/internal/handler/gateway_handler_warmup_intercept_unit_test.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,6 +153,7 @@ func newTestGatewayHandler(t *testing.T, group *service.Group, accounts []*servi
153153
nil, // deferredService
154154
nil, // claudeTokenProvider
155155
nil, // sessionLimitCache
156+
nil, // rpmCache
156157
nil, // digestStore
157158
)
158159

0 commit comments

Comments
 (0)