Skip to content

Commit e63c839

Browse files
author
QTom
committed
fix: address deep code review issues for RPM limiting
- Move IncrementRPM after Forward success to prevent phantom RPM consumption during account switch retries - Add base_rpm input sanitization (clamp to 0-10000) in Create/Update - Add WindowCost scheduling checks to legacy path sticky sessions (4 check sites + 4 prefetch sites), fixing pre-existing gap - Clean up rpm_strategy/rpm_sticky_buffer when disabling RPM in BulkEditModal (JSONB merge cannot delete keys, use empty values) - Add json.Number test cases to TestGetBaseRPM/TestGetRPMStickyBuffer - Document TOCTOU race as accepted soft-limit design trade-off
1 parent 4b72aa3 commit e63c839

5 files changed

Lines changed: 92 additions & 27 deletions

File tree

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

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -483,6 +483,8 @@ func (h *AccountHandler) Create(c *gin.Context) {
483483
response.BadRequest(c, "rate_multiplier must be >= 0")
484484
return
485485
}
486+
// base_rpm 输入校验：负值归零，超过 10000 截断
487+
sanitizeExtraBaseRPM(req.Extra)
486488

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

556560
// 确定是否跳过混合渠道检查
557561
skipCheck := req.ConfirmMixedChannelRisk != nil && *req.ConfirmMixedChannelRisk
@@ -1736,3 +1740,44 @@ func (h *AccountHandler) BatchRefreshTier(c *gin.Context) {
17361740
func (h *AccountHandler) GetAntigravityDefaultModelMapping(c *gin.Context) {
17371741
response.Success(c, domain.DefaultAntigravityModelMapping)
17381742
}
1743+
1744+
// sanitizeExtraBaseRPM 对 extra map 中的 base_rpm 值进行范围校验和归一化。
1745+
// 负值归零，超过 10000 截断为 10000。extra 为 nil 或不含 base_rpm 时无操作。
1746+
func sanitizeExtraBaseRPM(extra map[string]any) {
1747+
if extra == nil {
1748+
return
1749+
}
1750+
raw, ok := extra["base_rpm"]
1751+
if !ok {
1752+
return
1753+
}
1754+
v := parseExtraIntForValidation(raw)
1755+
if v < 0 {
1756+
v = 0
1757+
} else if v > 10000 {
1758+
v = 10000
1759+
}
1760+
extra["base_rpm"] = v
1761+
}
1762+
1763+
// parseExtraIntForValidation 从 extra 字段的 any 值解析为 int，用于输入校验。
1764+
// 支持 int, int64, float64, json.Number, string 类型。
1765+
func parseExtraIntForValidation(value any) int {
1766+
switch v := value.(type) {
1767+
case int:
1768+
return v
1769+
case int64:
1770+
return int(v)
1771+
case float64:
1772+
return int(v)
1773+
case json.Number:
1774+
if i, err := v.Int64(); err == nil {
1775+
return int(i)
1776+
}
1777+
case string:
1778+
if i, err := strconv.Atoi(strings.TrimSpace(v)); err == nil {
1779+
return i
1780+
}
1781+
}
1782+
return 0
1783+
}

‎backend/internal/handler/gateway_handler.go‎

Lines changed: 19 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -366,13 +366,6 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
366366
// 账号槽位/等待计数需要在超时或断开时安全回收
367367
accountReleaseFunc = wrapReleaseOnDone(c.Request.Context(), accountReleaseFunc)
368368

369-
// RPM 计数递增（调度成功后、Forward 前）
370-
if account.IsAnthropicOAuthOrSetupToken() && account.GetBaseRPM() > 0 {
371-
if err := h.gatewayService.IncrementAccountRPM(c.Request.Context(), account.ID); err != nil {
372-
reqLog.Warn("gateway.rpm_increment_failed", zap.Int64("account_id", account.ID), zap.Error(err))
373-
}
374-
}
375-
376369
// 转发请求 - 根据账号平台分流
377370
var result *service.ForwardResult
378371
requestCtx := c.Request.Context()
@@ -410,6 +403,15 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
410403
return
411404
}
412405

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+
413415
// 捕获请求信息（用于异步记录，避免在 goroutine 中访问 gin.Context）
414416
userAgent := c.GetHeader("User-Agent")
415417
clientIP := ip.GetClientIP(c)
@@ -556,13 +558,6 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
556558
// 账号槽位/等待计数需要在超时或断开时安全回收
557559
accountReleaseFunc = wrapReleaseOnDone(c.Request.Context(), accountReleaseFunc)
558560

559-
// RPM 计数递增（调度成功后、Forward 前）
560-
if account.IsAnthropicOAuthOrSetupToken() && account.GetBaseRPM() > 0 {
561-
if err := h.gatewayService.IncrementAccountRPM(c.Request.Context(), account.ID); err != nil {
562-
reqLog.Warn("gateway.rpm_increment_failed", zap.Int64("account_id", account.ID), zap.Error(err))
563-
}
564-
}
565-
566561
// 转发请求 - 根据账号平台分流
567562
var result *service.ForwardResult
568563
requestCtx := c.Request.Context()
@@ -609,7 +604,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
609604
h.handleStreamingAwareError(c, status, code, message, streamStarted)
610605
return
611606
}
612-
// 兜底重试按“直接请求兜底分组”处理：清除强制平台，允许按分组平台调度
607+
// 兜底重试按"直接请求兜底分组"处理：清除强制平台，允许按分组平台调度
613608
ctx := context.WithValue(c.Request.Context(), ctxkey.ForcePlatform, "")
614609
c.Request = c.Request.WithContext(ctx)
615610
currentAPIKey = fallbackAPIKey
@@ -643,6 +638,15 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
643638
return
644639
}
645640

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+
646650
// 捕获请求信息（用于异步记录，避免在 goroutine 中访问 gin.Context）
647651
userAgent := c.GetHeader("User-Agent")
648652
clientIP := ip.GetClientIP(c)

‎backend/internal/service/account_rpm_test.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
package service
22

3-
import "testing"
3+
import (
4+
"encoding/json"
5+
"testing"
6+
)
47

58
func TestGetBaseRPM(t *testing.T) {
69
tests := []struct {
@@ -16,6 +19,7 @@ func TestGetBaseRPM(t *testing.T) {
1619
{"string value", map[string]any{"base_rpm": "15"}, 15},
1720
{"negative value", map[string]any{"base_rpm": -5}, 0},
1821
{"int64 value", map[string]any{"base_rpm": int64(20)}, 20},
22+
{"json.Number value", map[string]any{"base_rpm": json.Number("25")}, 25},
1923
}
2024
for _, tt := range tests {
2125
t.Run(tt.name, func(t *testing.T) {
@@ -103,6 +107,7 @@ func TestGetRPMStickyBuffer(t *testing.T) {
103107
{"custom buffer=0 fallback to default", map[string]any{"base_rpm": 10, "rpm_sticky_buffer": 0}, 2},
104108
{"custom buffer negative fallback", map[string]any{"base_rpm": 10, "rpm_sticky_buffer": -1}, 2},
105109
{"custom buffer with float", map[string]any{"base_rpm": 10, "rpm_sticky_buffer": float64(7)}, 7},
110+
{"json.Number base_rpm", map[string]any{"base_rpm": json.Number("10")}, 2},
106111
}
107112
for _, tt := range tests {
108113
t.Run(tt.name, func(t *testing.T) {

‎backend/internal/service/gateway_service.go‎

Lines changed: 17 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2242,6 +2242,9 @@ func (s *GatewayService) isAccountSchedulableForRPM(ctx context.Context, account
22422242
}
22432243

22442244
// IncrementAccountRPM increments the RPM counter for the given account.
2245+
// 已知 TOCTOU 竞态：调度时读取 RPM 计数与此处递增之间存在时间窗口，
2246+
// 高并发下可能短暂超出 RPM 限制。这是与 WindowCost 一致的 soft-limit
2247+
// 设计权衡——可接受的少量超额优于加锁带来的延迟和复杂度。
22452248
func (s *GatewayService) IncrementAccountRPM(ctx context.Context, accountID int64) error {
22462249
if s.rpmCache == nil {
22472250
return nil
@@ -2444,7 +2447,7 @@ func sameAccountWithLoadGroup(a, b accountWithLoad) bool {
24442447
// shuffleWithinPriorityAndLastUsed 对排序后的 []*Account 切片，按 (Priority, LastUsedAt) 分组后组内随机打乱。
24452448
//
24462449
// 注意：当 preferOAuth=true 时，需要保证 OAuth 账号在同组内仍然优先，否则会把排序时的偏好打散掉。
2447-
// 因此这里采用“组内分区 + 分区内 shuffle”的方式：
2450+
// 因此这里采用"组内分区 + 分区内 shuffle"的方式：
24482451
// - 先把同组账号按 (OAuth / 非 OAuth) 拆成两段，保持 OAuth 段在前；
24492452
// - 再分别在各段内随机打散，避免热点。
24502453
func shuffleWithinPriorityAndLastUsed(accounts []*Account, preferOAuth bool) {
@@ -2584,7 +2587,7 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context,
25842587
if clearSticky {
25852588
_ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash)
25862589
}
2587-
if !clearSticky && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForRPM(ctx, account, true) {
2590+
if !clearSticky && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForWindowCost(ctx, account, true) && s.isAccountSchedulableForRPM(ctx, account, true) {
25882591
if s.debugModelRoutingEnabled() {
25892592
logger.LegacyPrintf("service.gateway", "[ModelRoutingDebug] legacy routed sticky hit: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), accountID)
25902593
}
@@ -2607,7 +2610,8 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context,
26072610
}
26082611
accountsLoaded = true
26092612

2610-
// 提前预取 RPM 计数，确保 routing 段内的 isAccountSchedulableForRPM 调用能命中缓存
2613+
// 提前预取窗口费用+RPM 计数，确保 routing 段内的调度检查调用能命中缓存
2614+
ctx = s.withWindowCostPrefetch(ctx, accounts)
26112615
ctx = s.withRPMPrefetch(ctx, accounts)
26122616

26132617
routingSet := make(map[int64]struct{}, len(routingAccountIDs))
@@ -2690,7 +2694,7 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context,
26902694
if clearSticky {
26912695
_ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash)
26922696
}
2693-
if !clearSticky && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForRPM(ctx, account, true) {
2697+
if !clearSticky && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForWindowCost(ctx, account, true) && s.isAccountSchedulableForRPM(ctx, account, true) {
26942698
return account, nil
26952699
}
26962700
}
@@ -2711,7 +2715,8 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context,
27112715
}
27122716
}
27132717

2714-
// 批量预取 RPM 计数，避免逐个账号查询（N+1）
2718+
// 批量预取窗口费用+RPM 计数，避免逐个账号查询（N+1）
2719+
ctx = s.withWindowCostPrefetch(ctx, accounts)
27152720
ctx = s.withRPMPrefetch(ctx, accounts)
27162721

27172722
// 3. 按优先级+最久未用选择（考虑模型支持）
@@ -2804,7 +2809,7 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g
28042809
if clearSticky {
28052810
_ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash)
28062811
}
2807-
if !clearSticky && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForRPM(ctx, account, true) {
2812+
if !clearSticky && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForWindowCost(ctx, account, true) && s.isAccountSchedulableForRPM(ctx, account, true) {
28082813
if account.Platform == nativePlatform || (account.Platform == PlatformAntigravity && account.IsMixedSchedulingEnabled()) {
28092814
if s.debugModelRoutingEnabled() {
28102815
logger.LegacyPrintf("service.gateway", "[ModelRoutingDebug] legacy mixed routed sticky hit: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), accountID)
@@ -2825,7 +2830,8 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g
28252830
}
28262831
accountsLoaded = true
28272832

2828-
// 提前预取 RPM 计数，确保 routing 段内的 isAccountSchedulableForRPM 调用能命中缓存
2833+
// 提前预取窗口费用+RPM 计数，确保 routing 段内的调度检查调用能命中缓存
2834+
ctx = s.withWindowCostPrefetch(ctx, accounts)
28292835
ctx = s.withRPMPrefetch(ctx, accounts)
28302836

28312837
routingSet := make(map[int64]struct{}, len(routingAccountIDs))
@@ -2912,7 +2918,7 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g
29122918
if clearSticky {
29132919
_ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash)
29142920
}
2915-
if !clearSticky && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForRPM(ctx, account, true) {
2921+
if !clearSticky && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForWindowCost(ctx, account, true) && s.isAccountSchedulableForRPM(ctx, account, true) {
29162922
if account.Platform == nativePlatform || (account.Platform == PlatformAntigravity && account.IsMixedSchedulingEnabled()) {
29172923
return account, nil
29182924
}
@@ -2931,7 +2937,8 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g
29312937
}
29322938
}
29332939

2934-
// 批量预取 RPM 计数，避免逐个账号查询（N+1）
2940+
// 批量预取窗口费用+RPM 计数，避免逐个账号查询（N+1）
2941+
ctx = s.withWindowCostPrefetch(ctx, accounts)
29352942
ctx = s.withRPMPrefetch(ctx, accounts)
29362943

29372944
// 3. 按优先级+最久未用选择（考虑模型支持和混合调度）
@@ -5304,7 +5311,7 @@ func (s *GatewayService) isThinkingBlockSignatureError(respBody []byte) bool {
53045311
}
53055312

53065313
func (s *GatewayService) shouldFailoverOn400(respBody []byte) bool {
5307-
// 只对“可能是兼容性差异导致”的 400 允许切换，避免无意义重试。
5314+
// 只对"可能是兼容性差异导致"的 400 允许切换，避免无意义重试。
53085315
// 默认保守：无法识别则不切换。
53095316
msg := strings.ToLower(strings.TrimSpace(extractUpstreamErrorMessage(respBody)))
53105317
if msg == "" {

‎frontend/src/components/account/BulkEditAccountModal.vue‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1224,8 +1224,12 @@ const buildUpdatePayload = (): Record<string, unknown> | null => {
12241224
extra.rpm_sticky_buffer = bulkRpmStickyBuffer.value
12251225
}
12261226
} else {
1227-
// 关闭 RPM 限制 - 设置 base_rpm 为 0
1227+
// 关闭 RPM 限制 - 设置 base_rpm 为 0，并用空值覆盖关联字段
1228+
// 后端使用 JSONB || merge 语义，不会删除已有 key，
1229+
// 所以必须显式发送空值来重置（后端读取时会 fallback 到默认值）
12281230
extra.base_rpm = 0
1231+
extra.rpm_strategy = ''
1232+
extra.rpm_sticky_buffer = 0
12291233
}
12301234
updates.extra = extra
12311235
}

0 commit comments

Comments
 (0)