feat: add critical rate limit ip whitelist - #3226
Closed
jingx8885 wants to merge 5397 commits into
Closed
Conversation
- Change ESCAPE character from '\' to '!' for compatibility with MySQL/PostgreSQL/SQLite - Adjust sanitization logic to escape '!' and '_' correctly, improving input validation for search queries
…d improved rate limiting
fix: /v1/chat/completions -> /v1/responses json_schema
将散落在多个文件中的预扣费/结算/退款逻辑抽象为统一的 BillingSession 生命周期管理： - 新增 BillingSettler 接口 (relay/common/billing.go) 避免循环引用 - 新增 FundingSource 接口 + WalletFunding / SubscriptionFunding 实现 (service/funding_source.go) - 新增 BillingSession 封装预扣/结算/退款原子操作 (service/billing_session.go) - 新增 SettleBilling 统一结算辅助函数，替换各 handler 中的 quotaDelta 模式 - 重写 PreConsumeBilling 为 BillingSession 工厂入口 - controller/relay.go 退款守卫改用 BillingSession.Refund() 修复的 Bug： - 令牌额度泄漏：PreConsumeTokenQuota 成功但 DecreaseUserQuota 失败时未回滚 - 订阅退款遗漏：FinalPreConsumedQuota=0 但 SubscriptionPreConsumed>0 时跳过退款 - 订阅多扣费：subConsume 强制为 1 但 FinalPreConsumedQuota 不同步 - 退款路径不统一：钱包/订阅退款逻辑现统一由 FundingSource.Refund 分派
- Settle 部分失败保护：新增 fundingSettled 标记，资金来源提交后 令牌调整失败不再导致 Refund 误退已结算的资金 - 订阅多扣费修复：trySubscription 传 subConsume 而非 preConsumedQuota 给 preConsume，保证三者（amount/preConsume/FinalPreConsumedQuota）一致 - 令牌回滚错误记录：preConsume 中 funding 失败时令牌回滚错误不再丢弃 - 移除钱包路径死代码：用户额度不足的 strings.Contains 匹配不可能命中 - WalletFunding.Refund 不重试：IncreaseUserQuota 非幂等，重试会多退
…e recharge card tabs - Defaulting to subscriptions when available and avoiding initial flash when no plans exist. - Adjust the wide-screen layout to place wallet and invite sections side by side, simplify the subscription header and controls, and add padding to prevent card borders from clipping. - Update related i18n strings by adding the new tab label and removing the obsolete subscription blurb.
…-when-no-plans ✨ refactor(wallet): Top-up layout to embed subscription plans into the recharge card tabs
refactor: 抽象统一计费会话 BillingSession
Add a lightweight active-subscription check to skip subscription pre-consume when none exist, reducing unnecessary transactions and locks. In the subscription UI, disable subscription-first options when no active plan is available, show the effective fallback to wallet with a clear notice, and distinguish “invalidated” from “expired” states. Update i18n strings across supported locales to reflect the new messages and status labels.
Aligns the error variable types in the subscription-first path so that quota fallback checks use the correct NewAPIError. This prevents build failures and preserves the intended wallet fallback when subscription pre-consume returns an insufficient quota error.
Routes quota alerts through a subscription-specific check when billing from subscriptions, preventing wallet-based thresholds from triggering false warnings. Updates the notification settings description and localization keys to clarify that both wallet and subscription balances are monitored.
🔔 feat: Add subscription-aware quota notifications and update UI copy
…-fallback ✨ chore: Improve subscription billing fallback and UI states
当上游为 AWS Bedrock 时，message_delta 的 usage 可能缺少 input_tokens、 cache_creation_input_tokens、cache_read_input_tokens 等字段，导致与原生 Anthropic 格式不一致。从 message_start 积累的 claudeInfo 中补全这些字段后 重新序列化，确保客户端收到一致的 usage 格式。
Modified the formatUserLogs function to include a startIdx parameter, allowing for more flexible log ID assignment. Updated calls to this function in GetLogByTokenId and GetUserLogs to pass the appropriate starting index.
feat: add Codex channel disclaimer (i18n, OpenAI terms)
feat: Force beta=true parameter for Anthropic channel
feat(oauth): implement custom OAuth provider
fix: Claude stream block index/type transitions
…tter clarity and functionality
Keep the model pricing editor wording aligned with the new price-based UI while exposing cache, image, and audio pricing in the marketplace so users can see the full configured pricing model.
Introduce a billing display mode feature allowing users to toggle between price and ratio views. Update relevant components and hooks to support this new functionality, ensuring consistent pricing information is displayed across the application.
Add siteDisplayType prop across various pricing components to conditionally render pricing information based on the selected display type. This update enhances the user experience by ensuring that pricing details are accurately represented according to the chosen display mode, particularly for token-based views.
…amOverrideEditorModal
为渠道参数覆盖可视化规则提供拖拽排序支持
feat: improve Gemini cache observability and Vertex custom URL support
…b3b03ba703796ea3 fix: kling risk fail return openAIVideo error
fix: add explicit docker-compose networks
…eader-append feat：support $keep_only_declared and deduped $append for header override
Contributor
WalkthroughThis PR introduces channel-level RPM rate limiting via new middleware, implements explicit caching for Gemini API requests with diagnostics tracking, restructures channel auto-test logic with decision-based evaluation, adds custom base URL support for Vertex AI, and extends channel affinity rules for Gemini models. Frontend UI updated to configure per-channel RPM limits. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Middleware as ChannelRequestRateLimit
participant Redis
participant Memory as In-Memory Limiter
participant Handler
Client->>Middleware: HTTP Request
Middleware->>Middleware: Extract Channel ID & Settings
alt RPM Limit > 0
Middleware->>Middleware: Determine Backend (Redis or Memory)
alt Redis Enabled
Middleware->>Redis: Get rateLimit:CRPM:channel:{id}:{minute}
Redis-->>Middleware: Counter Value
Middleware->>Redis: Increment & Set TTL(70s)
Redis-->>Middleware: Success
else In-Memory
Middleware->>Memory: Get Counter for Channel:Minute
Memory-->>Middleware: Counter Value
Middleware->>Memory: Increment
Memory-->>Middleware: Success
end
Middleware->>Middleware: Counter <= Limit?
alt Exceeded
Middleware-->>Client: 429 Too Many Requests
else Allowed
Middleware->>Handler: Next Handler
Handler-->>Client: 200 OK
end
else RPM Limit <= 0
Middleware->>Handler: Bypass (No Limit)
Handler-->>Client: 200 OK
end
sequenceDiagram
participant Client
participant Adaptor as GeminiAdaptor
participant Cache as ExplicitCache
participant Redis
participant Gemini as Upstream API
Client->>Adaptor: ConvertGeminiRequest
Adaptor->>Cache: prepareExplicitCache(ctx, relayInfo, request)
Cache->>Cache: buildExplicitCachePrefix(request)
Cache-->>Cache: Split into Prefix + Tail
Cache->>Cache: explicitCachePrefixMeta(relayInfo, prefix)
Cache-->>Cache: Token Count + Hash
Cache->>Cache: Compute Redis Key
Cache->>Redis: getExplicitCacheEntry(key)
Redis-->>Cache: Cache Hit or Miss
alt Cache Miss
Cache->>Gemini: createGeminiCachedContent(ctx, relayInfo, prefix)
Gemini-->>Cache: Cached Content Name
Cache->>Redis: putExplicitCacheEntry(key, entry, ttl)
Redis-->>Cache: Stored
else Cache Hit
Cache-->>Cache: Reuse Existing Entry
end
Cache->>Cache: applyExplicitCachedContent(request, cachedContent, tail)
Cache-->>Adaptor: Modified Request
Adaptor-->>Client: Updated Gemini Request
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 ast-grep (0.41.0)web/src/components/table/channels/modals/EditChannelModal.jsxThanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (4)
service/openaicompat/policy_test.go (1)
15-19: Test relies on implicit global state.This test assumes the global
ChatCompletionsToResponsesPolicyis disabled by default. If another test modifies global settings or if Go runs tests in parallel, results could become non-deterministic.Consider explicitly setting up the policy state before testing, or mocking the global settings, to improve test isolation and reliability.
💡 Suggested approach
func TestShouldChatCompletionsUseResponsesGlobal_NonCodexRespectsPolicy(t *testing.T) { + // Explicitly verify the policy is disabled, or use ShouldChatCompletionsUseResponsesPolicy + // with a known policy state instead of relying on global defaults. if ShouldChatCompletionsUseResponsesGlobal(0, constant.ChannelTypeOpenAI, "gpt-5.2") { t.Fatalf("expected non-codex channel to follow global policy default (disabled)") } }Alternatively, test
ShouldChatCompletionsUseResponsesPolicydirectly with an explicitChatCompletionsToResponsesPolicy{Enabled: false}to decouple from global state.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@service/openaicompat/policy_test.go` around lines 15 - 19, The test TestShouldChatCompletionsUseResponsesGlobal_NonCodexRespectsPolicy relies on implicit global state; update it to not depend on ChatCompletionsToResponsesPolicy's default by either explicitly setting the global ChatCompletionsToResponsesPolicy (e.g., assign a disabled policy and defer restoring the original) before calling ShouldChatCompletionsUseResponsesGlobal, or better, call ShouldChatCompletionsUseResponsesPolicy directly with an explicit ChatCompletionsToResponsesPolicy{Enabled:false} and the same channel/version parameters to assert the expected result; reference the test function name TestShouldChatCompletionsUseResponsesGlobal_NonCodexRespectsPolicy and the helper functions ShouldChatCompletionsUseResponsesGlobal and ShouldChatCompletionsUseResponsesPolicy when making the change.setting/operation_setting/channel_affinity_setting_test.go (1)
23-38: Test does not cover all KeySources.The test asserts
len(geminiRule.KeySources) < 5but the actual implementation has 6 KeySources (includingtoken_keyat index 5). Consider adding an assertion for the 6th KeySource or using an exact count check.♻️ Proposed enhancement to validate all KeySources
- if len(geminiRule.KeySources) < 5 { - t.Fatalf("expected multiple gemini key sources, got %d", len(geminiRule.KeySources)) + if len(geminiRule.KeySources) != 6 { + t.Fatalf("expected 6 gemini key sources, got %d", len(geminiRule.KeySources)) } assertKeySource := func(index int, typ string, path string, key string) { t.Helper() if geminiRule.KeySources[index].Type != typ || geminiRule.KeySources[index].Path != path || geminiRule.KeySources[index].Key != key { t.Fatalf("unexpected key source[%d]: %#v", index, geminiRule.KeySources[index]) } } assertKeySource(0, "gjson", "metadata.conversation_id", "") assertKeySource(1, "gjson", "metadata.thread_id", "") assertKeySource(2, "gjson", "metadata.session_id", "") assertKeySource(3, "gjson", "metadata.user_id", "") assertKeySource(4, "context_int", "", "token_id") + assertKeySource(5, "context_string", "", "token_key")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@setting/operation_setting/channel_affinity_setting_test.go` around lines 23 - 38, The test currently checks for at least 5 KeySources but the implementation has 6; update the check on geminiRule.KeySources to expect the exact length (6) and add a call to assertKeySource for the sixth entry (index 5) verifying its Type and Key (the token_key entry) using the existing assertKeySource helper and the geminiRule.KeySources symbol to locate the code; ensure the length assertion and the new assertKeySource(5, "<expectedType>", "<expectedPath>", "token_key") match the actual implementation.middleware/rate-limit.go (1)
107-141: Performance: Parse whitelist once at startup instead of on every request.The
os.Getenvcall and subsequent parsing of CIDR/IP entries occurs on every request to rate-limited endpoints. Given thatCriticalRateLimitis applied to high-traffic authentication and API routes, this adds unnecessary overhead. Consider parsing the whitelist once during initialization.♻️ Proposed refactor to cache the parsed whitelist
+var criticalRateLimitWhitelist []netip.Prefix +var criticalRateLimitWhitelistIPs []netip.Addr + +func init() { + whitelist := os.Getenv("CRITICAL_RATE_LIMIT_WHITELIST") + if whitelist == "" { + return + } + for _, entry := range strings.Split(whitelist, ",") { + entry = strings.TrimSpace(entry) + if entry == "" { + continue + } + if prefix, err := netip.ParsePrefix(entry); err == nil { + criticalRateLimitWhitelist = append(criticalRateLimitWhitelist, prefix) + continue + } + if ip, err := netip.ParseAddr(entry); err == nil { + criticalRateLimitWhitelistIPs = append(criticalRateLimitWhitelistIPs, ip) + } + } +} + func isWhitelistedForCriticalRateLimit(ip string) bool { if ip == "" { return false } - - whitelist := os.Getenv("CRITICAL_RATE_LIMIT_WHITELIST") - if whitelist == "" { + if len(criticalRateLimitWhitelist) == 0 && len(criticalRateLimitWhitelistIPs) == 0 { return false } clientIP, err := netip.ParseAddr(ip) if err != nil { return false } - for _, entry := range strings.Split(whitelist, ",") { - entry = strings.TrimSpace(entry) - if entry == "" { - continue - } - - if prefix, err := netip.ParsePrefix(entry); err == nil { - if prefix.Contains(clientIP) { - return true - } - continue - } - - if allowedIP, err := netip.ParseAddr(entry); err == nil && allowedIP == clientIP { + for _, prefix := range criticalRateLimitWhitelist { + if prefix.Contains(clientIP) { return true } } - + for _, allowedIP := range criticalRateLimitWhitelistIPs { + if allowedIP == clientIP { + return true + } + } return false }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@middleware/rate-limit.go` around lines 107 - 141, The isWhitelistedForCriticalRateLimit function currently calls os.Getenv and parses CIDR/IP entries on every request; instead, create a cached parsed whitelist initialized once (e.g., in an init function or via sync.Once) by a loader like loadCriticalRateLimitWhitelist that reads CRITICAL_RATE_LIMIT_WHITELIST, parses entries into netip.Prefix and netip.Addr collections, and stores them in package-level variables; then simplify isWhitelistedForCriticalRateLimit to only parse the incoming ip once and check against the cached prefixes/addrs (ensure the loader runs at startup or on first call with sync.Once for thread-safety).middleware/rate_limit_test.go (1)
70-103: Add a CIDR whitelist case here.This only verifies exact-IP matching. The new
netip.ParsePrefixbranch inmiddleware/rate-limit.gois still untested, so the advertisedIP/CIDR whitelistbehavior can regress silently. A case like198.51.100.0/24with request IP198.51.100.20would close that gap.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@middleware/rate_limit_test.go` around lines 70 - 103, Test only covers exact-IP whitelist matching; add a CIDR case to exercise the netip.ParsePrefix branch by extending TestCriticalRateLimitWhitelistBypassesLimit (or adding a sibling test) to set CRITICAL_RATE_LIMIT_WHITELIST to a CIDR like "198.51.100.0/24", send a request with newCriticalLimitRequest using source "198.51.100.20:3456" and assert HTTP 200 using the same newCriticalRateLimitRouter flow so the CIDR whitelist path in middleware/rate-limit.go is validated.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@controller/channel.go`:
- Around line 441-443: The check in controller.channel (ChannelRPMLimit) only
rejects negatives but lets huge positive values reach
common.InMemoryRateLimiter.Request which preallocates a slice with capacity
maxRequestNum and can OOM; add an upper-bound guard (e.g., clamp ChannelRPMLimit
to a safe MAX_CHANNEL_RPM constant or return an error if
channel.GetSetting().ChannelRPMLimit > MAX_CHANNEL_RPM) before using it so
Request never receives an unbounded value; update any validation code path that
reads ChannelRPMLimit (the conditional currently using
channel.GetSetting().ChannelRPMLimit) to enforce the new cap.
In `@middleware/channel-rate-limit.go`:
- Around line 68-75: Replace the separate INCR and EXPIRE calls on common.RDB
with a single atomic EVAL Lua script that does: local v = redis.call('INCR',
KEYS[1]); if v == 1 then redis.call('EXPIRE', KEYS[1], ARGV[1]) end; return v; —
call it via common.RDB.Eval(ctx, script, []string{key}, ttlSeconds).Handle the
returned integer as the count and propagate errors from Eval instead of using
Incr/Expire so the counter always has a TTL when created.
In `@relay/channel/gemini/adaptor.go`:
- Around line 44-45: The explicit-cache fields set by prepareExplicitCache
(which calls service.UpdateGeminiCacheDiagnostics) are being overwritten by
service.MarkGeminiCacheDiagnostics; swap the call order so
service.MarkGeminiCacheDiagnostics(...) runs before prepareExplicitCache(...)
(do the same swap for the other occurrence around lines 191-192) so the
explicit-cache diagnostics from
prepareExplicitCache/UpdateGeminiCacheDiagnostics are preserved and not rebuilt
away by MarkGeminiCacheDiagnostics.
In `@relay/channel/gemini/explicit_cache.go`:
- Around line 134-160: getExplicitCacheEntry is serving stale local entries
because putExplicitCacheEntry writes to explicitCacheLocal without any expiry
metadata; when Redis keys expire the local entry remains valid forever. Fix by
adding expiry tracking: extend explicitCacheEntry with an Expiry time.Time (or
wrap the stored value with an expiry timestamp), set entry.Expiry =
time.Now().Add(ttl) inside putExplicitCacheEntry before storing (and still set
Redis TTL), and update getExplicitCacheEntry to check entry.Expiry and treat
expired entries as missing (and optionally delete them from explicitCacheLocal).
Alternatively, if you prefer not to change the struct, schedule removal in
putExplicitCacheEntry with time.AfterFunc(ttl, func(){
explicitCacheLocal.Delete(key) }) so local and Redis expiry align.
In `@service/gemini_cache_diagnostics.go`:
- Around line 58-70: The diagnostics currently report only
GeminiImplicitCacheMinInputTokens (via min_input_tokens), but
prepareExplicitCache() actually uses the max of
GeminiImplicitCacheMinInputTokens(...) and ExplicitCacheMinInputTokens, so
MarkGeminiCacheDiagnostics() should compute and report the effective minimum the
explicit-cache path enforces; change the diagnostic calculation to compute
effectiveMin := max(GeminiImplicitCacheMinInputTokens(relayInfo.OriginModelName
or UpstreamModelName), ExplicitCacheMinInputTokens()) and emit effectiveMin as
"min_input_tokens" and use it when setting "eligible_for_implicit_cache" so
diagnostics match prepareExplicitCache() behavior (refer to
prepareExplicitCache, MarkGeminiCacheDiagnostics,
GeminiImplicitCacheMinInputTokens, ExplicitCacheMinInputTokens,
min_input_tokens, eligible_for_implicit_cache).
---
Nitpick comments:
In `@middleware/rate_limit_test.go`:
- Around line 70-103: Test only covers exact-IP whitelist matching; add a CIDR
case to exercise the netip.ParsePrefix branch by extending
TestCriticalRateLimitWhitelistBypassesLimit (or adding a sibling test) to set
CRITICAL_RATE_LIMIT_WHITELIST to a CIDR like "198.51.100.0/24", send a request
with newCriticalLimitRequest using source "198.51.100.20:3456" and assert HTTP
200 using the same newCriticalRateLimitRouter flow so the CIDR whitelist path in
middleware/rate-limit.go is validated.
In `@middleware/rate-limit.go`:
- Around line 107-141: The isWhitelistedForCriticalRateLimit function currently
calls os.Getenv and parses CIDR/IP entries on every request; instead, create a
cached parsed whitelist initialized once (e.g., in an init function or via
sync.Once) by a loader like loadCriticalRateLimitWhitelist that reads
CRITICAL_RATE_LIMIT_WHITELIST, parses entries into netip.Prefix and netip.Addr
collections, and stores them in package-level variables; then simplify
isWhitelistedForCriticalRateLimit to only parse the incoming ip once and check
against the cached prefixes/addrs (ensure the loader runs at startup or on first
call with sync.Once for thread-safety).
In `@service/openaicompat/policy_test.go`:
- Around line 15-19: The test
TestShouldChatCompletionsUseResponsesGlobal_NonCodexRespectsPolicy relies on
implicit global state; update it to not depend on
ChatCompletionsToResponsesPolicy's default by either explicitly setting the
global ChatCompletionsToResponsesPolicy (e.g., assign a disabled policy and
defer restoring the original) before calling
ShouldChatCompletionsUseResponsesGlobal, or better, call
ShouldChatCompletionsUseResponsesPolicy directly with an explicit
ChatCompletionsToResponsesPolicy{Enabled:false} and the same channel/version
parameters to assert the expected result; reference the test function name
TestShouldChatCompletionsUseResponsesGlobal_NonCodexRespectsPolicy and the
helper functions ShouldChatCompletionsUseResponsesGlobal and
ShouldChatCompletionsUseResponsesPolicy when making the change.
In `@setting/operation_setting/channel_affinity_setting_test.go`:
- Around line 23-38: The test currently checks for at least 5 KeySources but the
implementation has 6; update the check on geminiRule.KeySources to expect the
exact length (6) and add a call to assertKeySource for the sixth entry (index 5)
verifying its Type and Key (the token_key entry) using the existing
assertKeySource helper and the geminiRule.KeySources symbol to locate the code;
ensure the length assertion and the new assertKeySource(5, "<expectedType>",
"<expectedPath>", "token_key") match the actual implementation.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 2dd9054c-9c25-4fb5-a690-ab5711337c0d
📒 Files selected for processing (26)
.env.exampleconstant/context_key.gocontroller/channel-test.gocontroller/channel.gocontroller/channel_auto_test_decision_test.godto/channel_settings.gomiddleware/channel-rate-limit.gomiddleware/channel_rate_limit_test.gomiddleware/rate-limit.gomiddleware/rate_limit_test.gorelay/channel/gemini/adaptor.gorelay/channel/gemini/explicit_cache.gorelay/channel/gemini/explicit_cache_test.gorelay/channel/vertex/adaptor.gorelay/channel/vertex/adaptor_test.gorouter/relay-router.gorouter/video-router.goservice/gemini_cache_diagnostics.goservice/log_info_generate.goservice/log_info_generate_gemini_test.goservice/openaicompat/policy.goservice/openaicompat/policy_test.gosetting/model_setting/gemini.gosetting/operation_setting/channel_affinity_setting.gosetting/operation_setting/channel_affinity_setting_test.goweb/src/components/table/channels/modals/EditChannelModal.jsx
Comment on lines
+441
to
+443
| if channel.GetSetting().ChannelRPMLimit < 0 { | ||
| return fmt.Errorf("渠道 RPM 限制不能小于 0") | ||
| } |
Contributor
There was a problem hiding this comment.
Cap ChannelRPMLimit on the upper bound too.
Negative values are rejected, but very large positive values still flow into common.InMemoryRateLimiter.Request, which preallocates make([]int64, 0, maxRequestNum) for a new bucket. A bad admin value can OOM the process on the first request to that channel.
💡 Suggested guard
func validateChannel(channel *model.Channel, isAdd bool) error {
+ const maxChannelRPMLimit = 100000
// 校验 channel settings
if err := channel.ValidateSettings(); err != nil {
return fmt.Errorf("渠道额外设置[channel setting] 格式错误：%s", err.Error())
}
- if channel.GetSetting().ChannelRPMLimit < 0 {
- return fmt.Errorf("渠道 RPM 限制不能小于 0")
+ if limit := channel.GetSetting().ChannelRPMLimit; limit < 0 || limit > maxChannelRPMLimit {
+ return fmt.Errorf("渠道 RPM 限制必须在 0 到 %d 之间", maxChannelRPMLimit)
}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@controller/channel.go` around lines 441 - 443, The check in
controller.channel (ChannelRPMLimit) only rejects negatives but lets huge
positive values reach common.InMemoryRateLimiter.Request which preallocates a
slice with capacity maxRequestNum and can OOM; add an upper-bound guard (e.g.,
clamp ChannelRPMLimit to a safe MAX_CHANNEL_RPM constant or return an error if
channel.GetSetting().ChannelRPMLimit > MAX_CHANNEL_RPM) before using it so
Request never receives an unbounded value; update any validation code path that
reads ChannelRPMLimit (the conditional currently using
channel.GetSetting().ChannelRPMLimit) to enforce the new cap.
Comment on lines
+68
to
+75
| count, err := common.RDB.Incr(ctx, key).Result() | ||
| if err != nil { | ||
| return false, err | ||
| } | ||
| if count == 1 { | ||
| if err := common.RDB.Expire(ctx, key, 70*time.Second).Err(); err != nil { | ||
| return false, err | ||
| } |
Contributor
There was a problem hiding this comment.
Make Redis bucket creation atomic.
If INCR succeeds and Line 73 EXPIRE fails, the counter is left without a TTL. That channel can then stay rate-limited until the key is deleted manually. Please collapse the increment and expiry into one Lua script / transactional operation.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@middleware/channel-rate-limit.go` around lines 68 - 75, Replace the separate
INCR and EXPIRE calls on common.RDB with a single atomic EVAL Lua script that
does: local v = redis.call('INCR', KEYS[1]); if v == 1 then redis.call('EXPIRE',
KEYS[1], ARGV[1]) end; return v; — call it via common.RDB.Eval(ctx, script,
[]string{key}, ttlSeconds).Handle the returned integer as the count and
propagate errors from Eval instead of using Incr/Expire so the counter always
has a TTL when created.
Comment on lines
+44
to
+45
| prepareExplicitCache(c, info, request) | ||
| service.MarkGeminiCacheDiagnostics(c, info, request, info.GetEstimatePromptTokens()) |
Contributor
There was a problem hiding this comment.
Preserve explicit-cache diagnostics before writing the base Gemini diagnostics.
prepareExplicitCache() already records fields like cache_action, cache_reason, and prefix metadata through service.UpdateGeminiCacheDiagnostics(). Calling service.MarkGeminiCacheDiagnostics() afterwards rebuilds the map and drops those fields, so the admin diagnostics lose the most useful explicit-cache details.
Suggested fix
- prepareExplicitCache(c, info, request)
- service.MarkGeminiCacheDiagnostics(c, info, request, info.GetEstimatePromptTokens())
+ service.MarkGeminiCacheDiagnostics(c, info, request, info.GetEstimatePromptTokens())
+ prepareExplicitCache(c, info, request)- prepareExplicitCache(c, info, geminiRequest)
- service.MarkGeminiCacheDiagnostics(c, info, geminiRequest, info.GetEstimatePromptTokens())
+ service.MarkGeminiCacheDiagnostics(c, info, geminiRequest, info.GetEstimatePromptTokens())
+ prepareExplicitCache(c, info, geminiRequest)Also applies to: 191-192
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@relay/channel/gemini/adaptor.go` around lines 44 - 45, The explicit-cache
fields set by prepareExplicitCache (which calls
service.UpdateGeminiCacheDiagnostics) are being overwritten by
service.MarkGeminiCacheDiagnostics; swap the call order so
service.MarkGeminiCacheDiagnostics(...) runs before prepareExplicitCache(...)
(do the same swap for the other occurrence around lines 191-192) so the
explicit-cache diagnostics from
prepareExplicitCache/UpdateGeminiCacheDiagnostics are preserved and not rebuilt
away by MarkGeminiCacheDiagnostics.
Comment on lines
+134
to
+160
| func getExplicitCacheEntry(key string) (explicitCacheEntry, bool) { | ||
| if common.RedisEnabled && common.RDB != nil { | ||
| if raw, err := common.RedisGet(key); err == nil && raw != "" { | ||
| var entry explicitCacheEntry | ||
| if err := common.UnmarshalJsonStr(raw, &entry); err == nil && entry.Name != "" { | ||
| return entry, true | ||
| } | ||
| } | ||
| } | ||
| if raw, ok := explicitCacheLocal.Load(key); ok { | ||
| if entry, ok := raw.(explicitCacheEntry); ok && entry.Name != "" { | ||
| return entry, true | ||
| } | ||
| } | ||
| return explicitCacheEntry{}, false | ||
| } | ||
|
|
||
| func putExplicitCacheEntry(key string, entry explicitCacheEntry, ttl time.Duration) { | ||
| if entry.Name == "" { | ||
| return | ||
| } | ||
| if common.RedisEnabled && common.RDB != nil { | ||
| if data, err := common.Marshal(entry); err == nil { | ||
| _ = common.RedisSet(key, string(data), ttl) | ||
| } | ||
| } | ||
| explicitCacheLocal.Store(key, entry) |
Contributor
There was a problem hiding this comment.
Expire the local explicit-cache entries too.
putExplicitCacheEntry always writes to explicitCacheLocal, but getExplicitCacheEntry never checks age. Once the upstream cachedContent TTL has elapsed, this process can keep reusing the stale name forever, and with Redis enabled Line 143 will still return that stale local entry after the Redis key has expired.
💡 One way to keep local and Redis expiry aligned
type explicitCacheEntry struct {
- Name string `json:"name"`
+ Name string `json:"name"`
+ ExpiresAt int64 `json:"expires_at"`
}
func getExplicitCacheEntry(key string) (explicitCacheEntry, bool) {
+ now := time.Now().Unix()
if common.RedisEnabled && common.RDB != nil {
if raw, err := common.RedisGet(key); err == nil && raw != "" {
var entry explicitCacheEntry
if err := common.UnmarshalJsonStr(raw, &entry); err == nil && entry.Name != "" {
+ if entry.ExpiresAt > 0 && entry.ExpiresAt <= now {
+ return explicitCacheEntry{}, false
+ }
return entry, true
}
}
}
if raw, ok := explicitCacheLocal.Load(key); ok {
if entry, ok := raw.(explicitCacheEntry); ok && entry.Name != "" {
+ if entry.ExpiresAt > 0 && entry.ExpiresAt <= now {
+ explicitCacheLocal.Delete(key)
+ return explicitCacheEntry{}, false
+ }
return entry, true
}
}
return explicitCacheEntry{}, false
}
func putExplicitCacheEntry(key string, entry explicitCacheEntry, ttl time.Duration) {
if entry.Name == "" {
return
}
+ entry.ExpiresAt = time.Now().Add(ttl).Unix()
if common.RedisEnabled && common.RDB != nil {
if data, err := common.Marshal(entry); err == nil {
_ = common.RedisSet(key, string(data), ttl)
}
}
explicitCacheLocal.Store(key, entry)
}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@relay/channel/gemini/explicit_cache.go` around lines 134 - 160,
getExplicitCacheEntry is serving stale local entries because
putExplicitCacheEntry writes to explicitCacheLocal without any expiry metadata;
when Redis keys expire the local entry remains valid forever. Fix by adding
expiry tracking: extend explicitCacheEntry with an Expiry time.Time (or wrap the
stored value with an expiry timestamp), set entry.Expiry = time.Now().Add(ttl)
inside putExplicitCacheEntry before storing (and still set Redis TTL), and
update getExplicitCacheEntry to check entry.Expiry and treat expired entries as
missing (and optionally delete them from explicitCacheLocal). Alternatively, if
you prefer not to change the struct, schedule removal in putExplicitCacheEntry
with time.AfterFunc(ttl, func(){ explicitCacheLocal.Delete(key) }) so local and
Redis expiry align.
Comment on lines
+58
to
+70
| minInputTokens := GeminiImplicitCacheMinInputTokens(relayInfo.OriginModelName) | ||
| if minInputTokens == 0 { | ||
| minInputTokens = GeminiImplicitCacheMinInputTokens(relayInfo.UpstreamModelName) | ||
| } | ||
|
|
||
| diagnostics := map[string]any{ | ||
| "strategy": model_setting.GetGeminiSettings().CacheStrategy, | ||
| "channel_id": channelID, | ||
| "channel_type": channelType, | ||
| "upstream_base_url": baseURL, | ||
| "request_input_tokens": requestInputTokens, | ||
| "min_input_tokens": minInputTokens, | ||
| "eligible_for_implicit_cache": minInputTokens > 0 && requestInputTokens >= minInputTokens, |
Contributor
There was a problem hiding this comment.
Report the same minimum-token threshold that the explicit-cache path actually enforces.
prepareExplicitCache() compares the prefix length against max(GeminiImplicitCacheMinInputTokens(...), ExplicitCacheMinInputTokens), but MarkGeminiCacheDiagnostics() reports only the implicit model minimum here. If ops raises ExplicitCacheMinInputTokens, min_input_tokens and the eligibility flag become misleading even though the cache path behaved correctly.
Suggested fix
minInputTokens := GeminiImplicitCacheMinInputTokens(relayInfo.OriginModelName)
if minInputTokens == 0 {
minInputTokens = GeminiImplicitCacheMinInputTokens(relayInfo.UpstreamModelName)
}
+ if model_setting.GetGeminiSettings().CacheStrategy == model_setting.GeminiCacheStrategyExplicitCache &&
+ model_setting.GetGeminiSettings().ExplicitCacheMinInputTokens > minInputTokens {
+ minInputTokens = model_setting.GetGeminiSettings().ExplicitCacheMinInputTokens
+ }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| minInputTokens := GeminiImplicitCacheMinInputTokens(relayInfo.OriginModelName) | |
| if minInputTokens == 0 { | |
| minInputTokens = GeminiImplicitCacheMinInputTokens(relayInfo.UpstreamModelName) | |
| } | |
| diagnostics := map[string]any{ | |
| "strategy": model_setting.GetGeminiSettings().CacheStrategy, | |
| "channel_id": channelID, | |
| "channel_type": channelType, | |
| "upstream_base_url": baseURL, | |
| "request_input_tokens": requestInputTokens, | |
| "min_input_tokens": minInputTokens, | |
| "eligible_for_implicit_cache": minInputTokens > 0 && requestInputTokens >= minInputTokens, | |
| minInputTokens := GeminiImplicitCacheMinInputTokens(relayInfo.OriginModelName) | |
| if minInputTokens == 0 { | |
| minInputTokens = GeminiImplicitCacheMinInputTokens(relayInfo.UpstreamModelName) | |
| } | |
| if model_setting.GetGeminiSettings().CacheStrategy == model_setting.GeminiCacheStrategyExplicitCache && | |
| model_setting.GetGeminiSettings().ExplicitCacheMinInputTokens > minInputTokens { | |
| minInputTokens = model_setting.GetGeminiSettings().ExplicitCacheMinInputTokens | |
| } | |
| diagnostics := map[string]any{ | |
| "strategy": model_setting.GetGeminiSettings().CacheStrategy, | |
| "channel_id": channelID, | |
| "channel_type": channelType, | |
| "upstream_base_url": baseURL, | |
| "request_input_tokens": requestInputTokens, | |
| "min_input_tokens": minInputTokens, | |
| "eligible_for_implicit_cache": minInputTokens > 0 && requestInputTokens >= minInputTokens, |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@service/gemini_cache_diagnostics.go` around lines 58 - 70, The diagnostics
currently report only GeminiImplicitCacheMinInputTokens (via min_input_tokens),
but prepareExplicitCache() actually uses the max of
GeminiImplicitCacheMinInputTokens(...) and ExplicitCacheMinInputTokens, so
MarkGeminiCacheDiagnostics() should compute and report the effective minimum the
explicit-cache path enforces; change the diagnostic calculation to compute
effectiveMin := max(GeminiImplicitCacheMinInputTokens(relayInfo.OriginModelName
or UpstreamModelName), ExplicitCacheMinInputTokens()) and emit effectiveMin as
"min_input_tokens" and use it when setting "eligible_for_implicit_cache" so
diagnostics match prepareExplicitCache() behavior (refer to
prepareExplicitCache, MarkGeminiCacheDiagnostics,
GeminiImplicitCacheMinInputTokens, ExplicitCacheMinInputTokens,
min_input_tokens, eligible_for_implicit_cache).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
19 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Verification
Notes
Summary by CodeRabbit
New Features
Improvements