fix(rate-limit): unify lease mechanism for all periodic cost limits#674
fix(rate-limit): unify lease mechanism for all periodic cost limits#674
Conversation
- Migrate User daily quota from checkUserDailyCost to checkCostLimitsWithLease - Align key-quota.ts and keys.ts to use DB direct (sumKeyCostInTimeRange) - All periodic limits (5h/daily/weekly/monthly) now use lease mechanism - Total limits remain on 5-min Redis cache (no time window applicable) Closes #673 Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
📝 WalkthroughWalkthrough将成本计算从 RateLimitService.getCurrentCost 直接查询迁移至基于时间范围的 sumKeyCostInTimeRange 聚合方式，替换 checkCostLimits 为 checkCostLimitsWithLease，并修复 getTotalUsageForKey 中的 Warmup 数据排除不一致问题。 Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks 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 |
Summary of ChangesHello @ding113, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly refactors the rate-limiting system by standardizing all periodic cost limits to use a lease-based mechanism. This change improves consistency in how limits are applied and reported across different entities (Key, User, Provider) and various timeframes. It also ensures that all quota-related UI components fetch data from a unified, direct database source, resolving potential discrepancies caused by mixed data retrieval strategies. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
|
Note Unit test generation is an Early Access feature. Expect some limitations and changes as we gather feedback and continue to improve it. Generating unit tests... This may take up to 20 minutes. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request successfully unifies the periodic rate limit checks to use a consistent lease-based mechanism, which is a great improvement for maintainability and performance. The changes also align data sources for quota display, enhancing consistency across the UI. My review includes a couple of suggestions for further refactoring to improve code clarity and robustness.
|
|
||
| if (!dailyCheck.allowed) { | ||
| logger.warn(`[RateLimit] User daily limit exceeded: user=${user.id}, ${dailyCheck.reason}`); | ||
| const { currentUsage, limitValue } = parseLimitInfo(userDailyCheck.reason!); |
There was a problem hiding this comment.
Parsing values from a formatted string can be brittle. A future change to the error message format in checkCostLimitsWithLease could break this parsing logic.
For better robustness and maintainability, consider refactoring checkCostLimitsWithLease to return a structured object with currentUsage and limitValue fields, similar to how checkTotalCostLimit returns a current value. This would eliminate the need for string parsing here.
github-actions
bot
added
size/M
|
Caution The CodeRabbit agent's plans did not produce any file changes. |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review Summary
No significant issues identified in this PR. The changes correctly migrate User daily quota from checkUserDailyCost to checkCostLimitsWithLease and align admin interfaces (key-quota.ts, keys.ts) to use DB direct queries for consistency with my-usage.ts.
PR Size: M
- Lines changed: 694 (534 additions, 160 deletions)
- Files changed: 9
Review Coverage
- Logic and correctness - Clean
- Migration from
checkCostLimitstocheckCostLimitsWithLeaseis consistent parseLimitInfocorrectly handles both old and new error message formats- Time range calculations use appropriate functions (
getTimeRangeForPeriod,getTimeRangeForPeriodWithMode)
- Migration from
- Security (OWASP Top 10) - Clean
- Error handling - Clean
- Error message parsing has appropriate fallback behavior
- Logging is consistent with existing patterns
- Type safety - Clean
DailyResetModetype casts are consistent with existing codebase patterns
- Documentation accuracy - Clean
- Comments accurately describe the migration and new behavior
- Test documentation clearly explains the consistency fix
- Test coverage - Adequate
- 221 lines of new tests covering
parseLimitInfoformat parsing - Tests verify both
checkCostLimitsandcheckCostLimitsWithLeaseformats - Documentation tests explain the lease usage matrix
- 221 lines of new tests covering
- Code clarity - Good
- Consistent naming and structure across modified files
- Clear separation of Key vs User time range calculations
Notes
- The
EXCLUDE_WARMUP_CONDITIONaddition togetTotalUsageForKeyensures consistency with other statistics functions - All periodic cost limits (5h/daily/weekly/monthly) now use the lease mechanism uniformly
- Total limits remain on 5-min Redis cache (no time window applicable) as documented
Automated review by Claude AI
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unit/proxy/rate-limit-guard.test.ts (1)
3-41: 测试 mock 的 ERROR_CODES 缺少 RATE_LIMIT_5H_ROLLING_EXCEEDED当前实现已使用
RATE_LIMIT_5H_ROLLING_EXCEEDED，但 mock 中未定义，会把undefined传给getErrorMessageServer，降低用例对回归的敏感度。建议补齐常量。建议修改
ERROR_CODES: { RATE_LIMIT_TOTAL_EXCEEDED: "RATE_LIMIT_TOTAL_EXCEEDED", RATE_LIMIT_CONCURRENT_SESSIONS_EXCEEDED: "RATE_LIMIT_CONCURRENT_SESSIONS_EXCEEDED", RATE_LIMIT_RPM_EXCEEDED: "RATE_LIMIT_RPM_EXCEEDED", RATE_LIMIT_DAILY_QUOTA_EXCEEDED: "RATE_LIMIT_DAILY_QUOTA_EXCEEDED", RATE_LIMIT_5H_EXCEEDED: "RATE_LIMIT_5H_EXCEEDED", + RATE_LIMIT_5H_ROLLING_EXCEEDED: "RATE_LIMIT_5H_ROLLING_EXCEEDED", RATE_LIMIT_WEEKLY_EXCEEDED: "RATE_LIMIT_WEEKLY_EXCEEDED", RATE_LIMIT_MONTHLY_EXCEEDED: "RATE_LIMIT_MONTHLY_EXCEEDED", },
🧹 Nitpick comments (3)
tests/unit/proxy/provider-selector-total-limit.test.ts (1)
148-159: 测试断言与新的 Lease 机制保持一致。测试正确验证了
checkCostLimitsWithLease的调用参数，包含所有必要的限额字段。建议：第一个测试（
filterByLimits测试，第 40-101 行）未显式验证checkCostLimitsWithLease的调用。根据provider-selector.ts中的实现，filterByLimits也会调用此方法。可以考虑在该测试中添加类似的断言以确保覆盖完整。tests/unit/actions/my-usage-consistency.test.ts (1)
89-119: 建议：将文档测试转换为真实的集成测试。当前测试使用
vi.doMock但未实际调用被 mock 的模块，导致 mock 未被使用。如果目的是文档记录，现有方式可接受。但如果希望真正验证数据源一致性，建议：
- 实际导入并调用
getMyQuota或相关函数- 验证 mock 函数被正确调用
- 或者将此类文档移至 README/ADR 文档中
示例：将文档测试转为真实测试
it("should use sumKeyCostInTimeRange for Key quota (not RateLimitService.getCurrentCost)", async () => { - // This test documents the expected behavior: - // Key quota should use direct DB query (sumKeyCostInTimeRange) instead of Redis-first (getCurrentCost) - - // Mock the statistics module const sumKeyCostInTimeRangeMock = vi.fn(async () => 10.5); - const sumKeyTotalCostByIdMock = vi.fn(async () => 100.25); - const sumUserCostInTimeRangeMock = vi.fn(async () => 10.5); - const sumUserTotalCostMock = vi.fn(async () => 100.25); - - vi.doMock("@/repository/statistics", () => ({ - sumKeyCostInTimeRange: sumKeyCostInTimeRangeMock, - sumKeyTotalCostById: sumKeyTotalCostByIdMock, - sumUserCostInTimeRange: sumUserCostInTimeRangeMock, - sumUserTotalCost: sumUserTotalCostMock, - })); - - // Verify the function signatures match - expect(typeof sumKeyCostInTimeRangeMock).toBe("function"); - expect(typeof sumKeyTotalCostByIdMock).toBe("function"); - - // The test validates that: - // 1. Key 5h/daily/weekly/monthly uses sumKeyCostInTimeRange (DB direct) - // ... + vi.mock("@/repository/statistics", () => ({ + sumKeyCostInTimeRange: sumKeyCostInTimeRangeMock, + // ... other mocks + })); + + // Import after mocking + const { getMyQuota } = await import("@/actions/my-usage"); + + // Call and verify + await getMyQuota(/* test params */); + expect(sumKeyCostInTimeRangeMock).toHaveBeenCalled(); });src/actions/my-usage.ts (1)
237-292: 建议复用已计算的时间范围并避免调用已弃用的 sumUserCost目前 5h/weekly/monthly/total 仍走
sumUserCost，会二次计算时间范围并引入额外异步开销，也可能与 Key 侧窗口存在微小偏差。可以直接用已算好的range5h/rangeWeekly/rangeMonthly+sumUserCostInTimeRange，并直接调用sumUserTotalCost，同时去掉对已弃用函数的依赖。建议修改
- const { sumUserCostInTimeRange, sumKeyCostInTimeRange, sumKeyTotalCostById } = await import( - "@/repository/statistics" - ); + const { + sumUserCostInTimeRange, + sumUserTotalCost, + sumKeyCostInTimeRange, + sumKeyTotalCostById, + } = await import("@/repository/statistics");- sumUserCost(user.id, "5h"), + sumUserCostInTimeRange(user.id, range5h.startTime, range5h.endTime), sumUserCostInTimeRange(user.id, userDailyTimeRange.startTime, userDailyTimeRange.endTime), - sumUserCost(user.id, "weekly"), - sumUserCost(user.id, "monthly"), - sumUserCost(user.id, "total"), + sumUserCostInTimeRange(user.id, rangeWeekly.startTime, rangeWeekly.endTime), + sumUserCostInTimeRange(user.id, rangeMonthly.startTime, rangeMonthly.endTime), + sumUserTotalCost(user.id),
github-project-automation
bot
moved this from Backlog
to Done
in Claude Code Hub Roadmap
1 participant
Summary
checkUserDailyCosttocheckCostLimitsWithLeasekey-quota.tsandkeys.tsto use DB direct (sumKeyCostInTimeRange) for consistency withmy-usage.tsProblem
The rate limiting system had inconsistent data sources and mechanisms:
checkUserDailyCostwhile other periodic limits usedcheckCostLimitsWithLeasekey-quota.ts,keys.ts) used Redis-first queries whilemy-usage.tsused DB direct queriesgetTotalUsageForKeydid not exclude warmup requests, causing potential billing inconsistenciesRelated Issues:
getTotalUsageForKeySolution
checkCostLimitsWithLeasefor consistency with other periodic limitssumKeyCostInTimeRangeinstead ofRateLimitService.getCurrentCostEXCLUDE_WARMUP_CONDITIONtogetTotalUsageForKeyparseLimitInfoto handle bothcheckCostLimitsandcheckCostLimitsWithLeaseerror formatsLease Usage Matrix (After Migration)
Changes
Core Changes
rate-limit-guard.ts: Migrate User daily check fromcheckUserDailyCosttocheckCostLimitsWithLease(+89/-73)provider-selector.ts: Update Provider checks to usecheckCostLimitsWithLease(+3/-3)usage-logs.ts: AddEXCLUDE_WARMUP_CONDITIONtogetTotalUsageForKey(+7/-1)Data Source Alignment
key-quota.ts: Switch fromRateLimitService.getCurrentCosttosumKeyCostInTimeRange(+24/-11)keys.ts: Switch fromRateLimitService.getCurrentCosttosumKeyCostInTimeRange(+25/-15)my-usage.ts: Align Key quota queries with User quota (DB direct) (+28/-19)Tests
my-usage-consistency.test.ts: New test file documenting data source consistency (+221 lines)rate-limit-guard.test.ts: Update mocks and add User daily rolling mode tests (+124/-29)provider-selector-total-limit.test.ts: Update mock method names (+13/-9)Test Plan
bun run typecheckpassesbun run lintpasses (modified files)Checklist
Description enhanced by Claude AI
Greptile Overview
Greptile Summary
Unified all periodic cost limits (5h/daily/weekly/monthly) to use the lease mechanism (
checkCostLimitsWithLease), migrating User daily quota from the standalonecheckUserDailyCostfunction.Key Changes:
checkCostLimitsWithLeasefor consistency with Key and Provider periodic limitsmy-usage.ts,key-quota.ts,keys.ts) now use direct DB queries (sumKeyCostInTimeRange) instead of Redis-first approach for data consistencyparseLimitInfoinrate-limit-guard.tsupdated to parse both old（X/Y）and new(usage: X/Y)error formatsgetTotalUsageForKeynow excludes warmup requests viaEXCLUDE_WARMUP_CONDITIONfor consistencyBenefits:
Confidence Score: 5/5
Important Files Changed
checkUserDailyCosttocheckCostLimitsWithLease, updatedparseLimitInfoto support both error message formatsEXCLUDE_WARMUP_CONDITIONtogetTotalUsageForKeyfor consistency with other statistics functionssumKeyCostInTimeRange,sumKeyTotalCostById) for data source consistencycheckCostLimitsWithLeaseformat, added new tests for User daily rolling mode and check order verificationSequence Diagram
sequenceDiagram participant Client participant RateLimitGuard participant RateLimitService participant LeaseCache participant Redis participant Database Client->>RateLimitGuard: API Request Note over RateLimitGuard: Check Order:<br/>1. Total limits<br/>2. Session/RPM<br/>3-6. Periodic limits (5h, daily)<br/>7-10. Weekly, Monthly RateLimitGuard->>RateLimitService: checkCostLimitsWithLease(key, "key", {5h, daily, weekly, monthly}) RateLimitService->>LeaseCache: Get cached lease slice alt Lease exists and valid LeaseCache-->>RateLimitService: Return cached slice Note over RateLimitService: Use cached budget<br/>(reduce DB pressure) else Lease expired or missing LeaseCache-->>RateLimitService: Cache miss RateLimitService->>Database: sumKeyCostInTimeRange(keyId, startTime, endTime) Database-->>RateLimitService: Current usage RateLimitService->>Redis: Store new lease slice (Lua script) Redis-->>RateLimitService: Lease created end RateLimitService->>RateLimitService: Check usage < limit RateLimitService-->>RateLimitGuard: {allowed: true/false, reason} alt User daily check needed RateLimitGuard->>RateLimitService: checkCostLimitsWithLease(user, "user", {daily}) Note over RateLimitService: Same lease mechanism<br/>for User quota RateLimitService-->>RateLimitGuard: {allowed: true/false, reason} end alt All checks pass RateLimitGuard-->>Client: Request allowed else Any check fails RateLimitGuard->>RateLimitGuard: parseLimitInfo(reason) Note over RateLimitGuard: Extract usage from:<br/>(usage: X/Y) format RateLimitGuard-->>Client: RateLimitError with details end